[Oryx] Add inverse rule for sqrt

PiperOrigin-RevId: 376972734

Author: sharadmv

spinoffs/oryx/oryx/core/interpreters/inverse/inverse_test.py

@@ -353,6 +353,18 @@ def f(x, y):
     onp.testing.assert_allclose(
         f_y_ildj(3.), np.log(np.abs(jax.grad(f_y_inv)(3.))))
 
+  def test_sqrt_inverse(self):
+    def f(x):
+      return np.sqrt(x)
+    f_inv = core.inverse(f)
+    onp.testing.assert_allclose(f_inv(2.), 4.)
+
+  def test_sqrt_ildj(self):
+    def f(x):
+      return np.sqrt(x)
+    f_ildj = core.ildj(f)
+    onp.testing.assert_allclose(f_ildj(3.), np.log(2.) + np.log(3.))
+
 
 if __name__ == '__main__':
   os.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=2'

spinoffs/oryx/oryx/core/interpreters/inverse/rules.py

@@ -48,6 +48,7 @@
 register_elementwise(lax.expm1_p)(np.log1p)
 register_elementwise(lax.log1p_p)(np.expm1)
 register_elementwise(lax.neg_p)(lambda x: -x)
+register_elementwise(lax.sqrt_p)(np.square)
 
 
 @register_elementwise(lax.integer_pow_p)