Really skip exp2 in Pallas GPU tests with older jaxlib

superbobry · jax authors · commit 2d65571b2155 · 2024-03-22T05:42:38.000-07:00
PiperOrigin-RevId: 618149873
diff --git a/tests/pallas/pallas_test.py b/tests/pallas/pallas_test.py
@@ -1551,7 +1551,7 @@ class PallasOpsTest(PallasTest):
       for fn, dtype in itertools.product(*args)
   )
   def test_elementwise(self, fn, dtype):
-    if fn is lax.exp2 and jaxlib_version < (0, 4, 26):
+    if fn is jnp.exp2 and jaxlib_version < (0, 4, 26):
       self.skipTest("Requires jaxlib 0.4.26 or later")
 
     @functools.partial(

Original file line number	Diff line number	Diff line change
`@@ -1551,7 +1551,7 @@ class PallasOpsTest(PallasTest):`
`1551`	`1551`	`for fn, dtype in itertools.product(*args)`
`1552`	`1552`	`)`
`1553`	`1553`	`def test_elementwise(self, fn, dtype):`
`1554`		`- if fn is lax.exp2 and jaxlib_version < (0, 4, 26):`
	`1554`	`+ if fn is jnp.exp2 and jaxlib_version < (0, 4, 26):`
`1555`	`1555`	`self.skipTest("Requires jaxlib 0.4.26 or later")`
`1556`	`1556`
`1557`	`1557`	`@functools.partial(`