Add sqrt/square primitives to the inverse registry after openxla/stab…

…lehlo#2623 PiperOrigin-RevId: 699910407
google-deepmind · Nov 25, 2024 · 0f8fe14 · 0f8fe14
1 parent 527d698
commit 0f8fe14
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/distrax/_src/utils/transformations.py b/distrax/_src/utils/transformations.py
@@ -93,6 +93,12 @@
     jax.lax.integer_pow_p: lambda x, y: jax.lax.pow_p.bind(x, 1.0/y)
 }
 
+if hasattr(jax.lax, "square_p"):
+  _inverse_registry.update({
+      jax.lax.square_p: jax.lax.sqrt_p,
+      jax.lax.sqrt_p: jax.lax.square_p,
+      jax.lax.rsqrt_p: lambda x: 1.0 / jax.lax.square_p.bind(x),
+  })
 
 _potentially_unstable_primitives = {
     jax.lax.tanh_p: "distrax.Tanh or distrax.Inverse(distrax.Tanh)",