def testThreefry2x32Empty(self): # Regression test for an op-by-op crash for empty arrays in CUDA mode. with api.disable_jit(): result = random.threefry_2x32( (np.uint32(0x13198a2e), np.uint32(0x03707344)), jnp.ones((10, 0,), jnp.uint32)) np.testing.assert_equal(result, np.zeros((10, 0,), dtype=np.uint32))
def test_disable_jit(self): effects = [] @api.jit def f(x): effects.append(1) return x with api.disable_jit(): f(2) f(2) assert len(effects) == 2 f(2) f(2) assert len(effects) == 3