mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #15953 from jakevdp:keyarray-dynamic-slice

Browse Source 
PiperOrigin-RevId: 530936580
This commit is contained in:
jax authors 2023-05-10 09:58:49 -07:00
commit d6828c9c35
2 changed files with 14 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
jax/_src/prng.py
View File

@ -482,7 +482,8 @@ class KeyTyRules:
@staticmethod
def dynamic_slice_mlir(ctx, aval_out, x, start_indices) -> ir.Value:
dtype = dtypes.canonicalize_dtype(np.dtype('int64'))
index_avals = ctx.avals_in[1:]
dtype = dtypes.canonicalize_dtype(index_avals[0].dtype if index_avals else 'int64')
key_shape = aval_out.dtype.impl.key_shape
trailing_zeros = [mlir.ir_constant(np.array(0, dtype))] * len(key_shape)
start_indices = (*start_indices, *trailing_zeros)
@ -492,7 +493,8 @@ class KeyTyRules:
@staticmethod
def dynamic_update_slice_mlir(ctx, aval_out, x, update, *start_indices) -> ir.Value:
dtype = dtypes.canonicalize_dtype(np.dtype('int64'))
index_avals = ctx.avals_in[2:]
dtype = dtypes.canonicalize_dtype(index_avals[0].dtype if index_avals else 'int64')
key_shape = aval_out.dtype.impl.key_shape
zeros = [mlir.ir_constant(np.array(0, dtype=dtype))] * len(key_shape)
start_indices = (*start_indices, *zeros)

11
tests/random_test.py
View File

@ -1801,10 +1801,19 @@ class KeyArrayTest(jtu.JaxTestCase):
def test_dynamic_slice(self):
ks = self.make_keys(3, 4)
ys = jax.jit(lambda x, i: lax.dynamic_slice_in_dim(x, i, 2))(ks, 1)
index = np.int16(1) # non-default int type to catch type errors.
ys = jax.jit(partial(lax.dynamic_slice_in_dim, slice_size=2))(ks, index)
self.assertIsInstance(ys, random.KeyArray)
self.assertEqual(ys.shape, (2, 4))
def test_dynamic_update_slice(self):
ks = self.make_keys(3, 4)
k = self.make_keys(1, 4)
index = np.int16(1) # non-default int type to catch type errors.
ys = jax.jit(partial(lax.dynamic_update_slice_in_dim, axis=0))(ks, k, index)
self.assertIsInstance(ys, random.KeyArray)
self.assertEqual(ys.shape, (3, 4))
def test_transpose(self):
ks = self.make_keys(3, 4)
ys = jax.jit(lambda x: x.T)(ks)