[Oryx] Add inverse rules for integer_pow and pow

PiperOrigin-RevId: 376048209

Author: sharadmv

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

@@ -257,9 +257,9 @@ def ildj_rule(incells, outcells, **params):
       incell, = incells
       if not incell.top() and outcell.top():
         val = outcell.val
-        f_sum = lambda x: f(x).sum()
-        ildj_ = outcell.ildj + np.log(jax.grad(f_sum)(val))
-        ndslice = NDSlice.new(f(val), ildj_)
+        f_sum = lambda x: f(x, **params).sum()
+        ildj_ = outcell.ildj + np.log(np.abs(jax.grad(f_sum)(val)))
+        ndslice = NDSlice.new(f(val, **params), ildj_)
         incells = [InverseAndILDJ(outcell.aval, [ndslice])]
       elif not outcell.top() and incell.top():
         outcells = [InverseAndILDJ.new(prim.bind(incell.val, **params))]

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

@@ -17,6 +17,7 @@
 
 from absl.testing import absltest
 import jax
+from jax import lax
 import jax.numpy as np
 import numpy as onp
 
@@ -296,6 +297,62 @@ def naive_logit(x):
                                 tfb.Sigmoid().forward_log_det_jacobian(
                                     -100., 0))
 
+  def test_integer_pow_inverse(self):
+    def f(x):
+      return lax.integer_pow(x, 2)
+    f_inv = core.inverse(f)
+    onp.testing.assert_allclose(f_inv(2.), np.sqrt(2.))
+    def f2(x):
+      return lax.integer_pow(x, 3)
+    f2_inv = core.inverse(f2)
+    onp.testing.assert_allclose(f2_inv(2.), onp.cbrt(2.))
+
+  def test_integer_pow_ildj(self):
+    def f(x):
+      return lax.integer_pow(x, 2)
+    f_ildj = core.ildj(f)
+    onp.testing.assert_allclose(
+        f_ildj(2.), tfb.Power(2.).inverse_log_det_jacobian(2.))
+    def f2(x):
+      return lax.integer_pow(x, 3)
+    f2_ildj = core.ildj(f2)
+    onp.testing.assert_allclose(
+        f2_ildj(2.), tfb.Power(3.).inverse_log_det_jacobian(2.))
+
+  def test_reciprocal_inverse(self):
+    def f(x):
+      return np.reciprocal(x)
+    f_inv = core.inverse(f)
+    onp.testing.assert_allclose(f_inv(2.), 0.5)
+
+  def test_reciprocal_ildj(self):
+    def f(x):
+      return np.reciprocal(x)
+    f_ildj = core.ildj(f)
+    onp.testing.assert_allclose(f_ildj(2.), onp.log(1 / 4.))
+
+  def test_pow_inverse(self):
+    def f(x, y):
+      return lax.pow(x, y)
+    f_x_inv = core.inverse(lambda x: f(x, 2.))
+    onp.testing.assert_allclose(f_x_inv(2.), np.sqrt(2.))
+    f_y_inv = core.inverse(lambda y: f(2., y))
+    onp.testing.assert_allclose(f_y_inv(3.), np.log(3.) / np.log(2.))
+
+  def test_pow_ildj(self):
+    def f(x, y):
+      return lax.pow(x, y)
+    f_x_ildj = core.ildj(lambda x: f(x, 2.))
+    onp.testing.assert_allclose(
+        f_x_ildj(3.), tfb.Power(2.).inverse_log_det_jacobian(3.))
+    f_y_ildj = core.ildj(lambda y: f(2., y))
+    f_y_inv = core.inverse(lambda y: f(2., y))
+    y = f_y_inv(3.)
+    onp.testing.assert_allclose(
+        f_y_ildj(3.), -np.log(np.abs(jax.grad(lambda y: f(2., y))(y))))
+    onp.testing.assert_allclose(
+        f_y_ildj(3.), np.log(np.abs(jax.grad(f_y_inv)(3.))))
+
 
 if __name__ == '__main__':
   os.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=2'

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

@@ -50,6 +50,38 @@
 register_elementwise(lax.neg_p)(lambda x: -x)
 
 
+@register_elementwise(lax.integer_pow_p)
+def integer_pow_inverse(z, *, y):
+  """Inverse for `integer_pow_p` primitive."""
+  if y == 0:
+    raise ValueError('Cannot invert raising to a value to the 0-th power.')
+  elif y == 1:
+    return z
+  elif y == -1:
+    return np.reciprocal(z)
+  elif y == 2:
+    return np.sqrt(z)
+  return lax.pow(z, 1. / y)
+
+
+def pow_left(x, z, ildj_):
+  # x ** y = z
+  # y = f^-1(z) = log(z) / log(x)
+  # grad(f^-1)(z) = 1 / (z log(x))
+  # log(grad(f^-1)(z)) = log(1 / (z log(x))) = -log(z) - log(log(x))
+  return np.log(z) / np.log(x), ildj_ - np.log(z) - np.log(np.log(x))
+
+
+def pow_right(y, z, ildj_):
+  # x ** y = z
+  # x = f^-1(z) = z ** (1 / y)
+  # grad(f^-1)(z) = 1 / y * z ** (1 / y - 1)
+  # log(grad(f^-1)(z)) = (1 / y - 1)log(z) - log(y)
+  y_inv = np.reciprocal(y)
+  return lax.pow(z, y_inv), ildj_ + (y_inv - 1.) * np.log(z) - np.log(y)
+register_binary(lax.pow_p)(pow_left, pow_right)
+
+
 def add_left(left_val, out_val, ildj_):
   return out_val - left_val, ildj_
 
@@ -182,15 +214,6 @@ def expit_ildj(y):
 def logit_ildj(y):
   return -jax.nn.softplus(-y) - jax.nn.softplus(y)
 
-# Monkey patching JAX so we can define custom, more numerically stable inverses.
-jax.scipy.special.expit = custom_inverse(jax.scipy.special.expit)
-jax.scipy.special.logit = custom_inverse(jax.scipy.special.logit)
-jax.nn.sigmoid = jax.scipy.special.expit
-jax.scipy.special.expit.def_inverse_unary(f_inv=jax.scipy.special.logit,
-                                          f_ildj=expit_ildj)
-jax.scipy.special.logit.def_inverse_unary(f_inv=jax.scipy.special.expit,
-                                          f_ildj=logit_ildj)
-
 
 def convert_element_type_ildj(incells, outcells, *, new_dtype, **params):
   """InverseAndILDJ rule for convert_element_type primitive."""
@@ -206,3 +229,14 @@ def convert_element_type_ildj(incells, outcells, *, new_dtype, **params):
             val, new_dtype=incell.aval.dtype, **params))]
   return incells, outcells, None
 ildj_registry[lax.convert_element_type_p] = convert_element_type_ildj
+
+
+# Monkey patching JAX so we can define custom, more numerically stable inverses.
+jax.scipy.special.expit = custom_inverse(jax.scipy.special.expit)
+jax.scipy.special.logit = custom_inverse(jax.scipy.special.logit)
+jax.nn.sigmoid = jax.scipy.special.expit
+jax.scipy.special.expit.def_inverse_unary(f_inv=jax.scipy.special.logit,
+                                          f_ildj=expit_ildj)
+jax.scipy.special.logit.def_inverse_unary(f_inv=jax.scipy.special.expit,
+                                          f_ildj=logit_ildj)
+