TMP

jax-ml · Mar 12, 2021 · 87e4a66 · 87e4a66
1 parent c1a8a19
commit 87e4a66
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/x64_context_test.py b/tests/x64_context_test.py
@@ -163,10 +163,10 @@ def test_convert_element_type(self):
       x = jnp.int64(1)
     self.assertEqual(x.dtype, jnp.int64)
 
-    y = lax.convert_element_type(x, jnp.int32)
+    y = x.astype(jnp.int32)
     self.assertEqual(y.dtype, jnp.int32)
 
-    z = api.jit(lax.convert_element_type, static_argnums=1)(x, jnp.int32)
+    z = api.jit(lambda x: x.astype(jnp.int32))(x)
     self.assertEqual(z.dtype, jnp.int32)
 
 if __name__ == "__main__":