feat: select (#34)

* feat: select * feat: numpy.power Signed-off-by: nstarman <nstarman@users.noreply.github.com>
GalacticDynamics · Mar 20, 2024 · 4ccca52 · 4ccca52
1 parent 0591ebc
commit 4ccca52
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/src/quaxed/lax.py b/src/quaxed/lax.py
@@ -12,8 +12,10 @@
     "erf",
     "erfc",
     "erf_inv",
+    "select",
 ]
 
+
 from jax import lax
 from quax import quaxify
 
@@ -28,3 +30,4 @@
 erf = quaxify(lax.erf)
 erfc = quaxify(lax.erfc)
 erf_inv = quaxify(lax.erf_inv)
+select = quaxify(lax.select)
diff --git a/src/quaxed/numpy/_core.py b/src/quaxed/numpy/_core.py
@@ -14,6 +14,7 @@
     "isclose",
     "matmul",
     "moveaxis",
+    "power",
     "squeeze",
     "trace",
     "vectorize",
@@ -68,6 +69,7 @@ def transfer_doc(func: T) -> T:
 isclose = quaxify(jnp.isclose)
 matmul = quaxify(jnp.matmul)
 moveaxis = quaxify(jnp.moveaxis)
+power = quaxify(jnp.power)
 squeeze = quaxify(jnp.squeeze)
 trace = quaxify(jnp.trace)