# Copyright 2024 The JAX Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # https://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import functools import math import sys from absl.testing import absltest from absl.testing import parameterized import jax from jax import lax from jax._src import dtypes from jax._src import test_util as jtu from jax._src import hypothesis_test_util as htu from jax._src.pallas import pallas_test_util as ptu from jax.experimental import pallas as pl import jax.numpy as jnp import numpy as np if sys.platform != "win32": from jax.experimental.pallas import tpu as pltpu else: pltpu = None jax.config.parse_flags_with_absl() htu.setup_hypothesis(max_examples=100) _JAX_DTYPES_NO_BOOL = ( jnp.float32, jnp.bfloat16, jnp.int32, jnp.int16, jnp.int8, jnp.int4, jnp.float8_e5m2, ) _JAX_DTYPES = ( *_JAX_DTYPES_NO_BOOL, jnp.bool_, ) _JAX_INT_DTYPES = ( jnp.int32, jnp.int16, jnp.int8, jnp.int4, jnp.uint32, jnp.uint16, jnp.uint8, jnp.uint4, ) def rand( shape: tuple[int, ...], dtype: np.dtype | jnp.dtype, seed: int = 1234 ) -> np.ndarray: """A helper function to generate random data for testing.""" rng = np.random.Generator(np.random.Philox(counter=0, key=seed)) if jnp.issubdtype(dtype, jnp.floating): return rng.normal(size=shape).astype(dtype) if jnp.issubdtype(dtype, jnp.integer): return rng.integers( jnp.iinfo(dtype).min, jnp.iinfo(dtype).max, shape, dtype=np.int32 ).astype(dtype) raise NotImplementedError(f"Unsupported random data generation for {dtype=}") @jtu.thread_unsafe_test_class(condition=not htu.hypothesis_is_thread_safe()) class OpsTest(ptu.PallasTPUTest): @parameterized.product( from_dtype=_JAX_DTYPES, to_dtype=_JAX_DTYPES, is_ref_bitcast=[False, True], use_primitive_io_op=[False, True], ) def test_bitcast( self, from_dtype, to_dtype, is_ref_bitcast, use_primitive_io_op ): if not jtu.is_device_tpu_at_least(version=4): self.skipTest("Run on TPUv4+ to have expected memory layout") if from_dtype == to_dtype: self.skipTest("No bitcast needed") if from_dtype == jnp.bool_ or to_dtype == jnp.bool_: self.skipTest("Bitcasting with bool is not supported") def kernel(x_ref, y_ref): if is_ref_bitcast: if use_primitive_io_op: pltpu.store(y_ref, pltpu.load(x_ref.bitcast(to_dtype))) else: y_ref[...] = x_ref.bitcast(to_dtype)[...] else: if use_primitive_io_op: pltpu.store(y_ref, pltpu.bitcast(pltpu.load(x_ref), to_dtype)) else: y_ref[...] = pltpu.bitcast(x_ref[...], to_dtype) m, n = 1, 256 in_packing = 32 // dtypes.itemsize_bits(from_dtype) out_packing = 32 // dtypes.itemsize_bits(to_dtype) in_shape = (m * in_packing, n) out_shape = (m * out_packing, n) inp = np.arange(np.prod(in_shape), dtype=from_dtype).reshape(in_shape) out = self.pallas_call( kernel, out_shape=jax.ShapeDtypeStruct(out_shape, to_dtype), )(inp) if not self.INTERPRET: out_interpret = pl.pallas_call( kernel, out_shape=jax.ShapeDtypeStruct(out_shape, to_dtype), interpret=True, )(inp) self.assertAllClose(out, out_interpret) @parameterized.product( from_dtype=_JAX_DTYPES, to_dtype=_JAX_DTYPES, ) def test_bitcast_scalar(self, from_dtype, to_dtype): if from_dtype == to_dtype: self.skipTest("No bitcast needed") if dtypes.itemsize_bits(from_dtype) != dtypes.itemsize_bits(to_dtype): self.skipTest("Bitwidths must match for scalar bitcast") if from_dtype == jnp.bool_ or to_dtype == jnp.bool_: self.skipTest("Bitcasting with bool is not supported") @functools.partial( self.pallas_call, out_shape=jax.ShapeDtypeStruct((1,), to_dtype), in_specs=[pl.BlockSpec(memory_space=pltpu.SMEM)], out_specs=pl.BlockSpec(memory_space=pltpu.SMEM), ) def _kernel(x_ref, o_ref): o_ref[0] = jax.lax.bitcast_convert_type(x_ref[0], to_dtype) x = jnp.array([1], dtype=from_dtype) out = _kernel(x) expected = jax.lax.bitcast_convert_type(x, to_dtype) self.assertAllClose(out, expected) def test_stop_gradient(self): def kernel(x_ref, y_ref): y_ref[...] = jax.lax.stop_gradient(x_ref[...] + 1) x = jnp.arange(1024, dtype=jnp.float32) y = pl.pallas_call(kernel, out_shape=x)(x) self.assertAllClose(y, x + 1) def test_interleave_vectors(self): def kernel(x_ref, y_ref, out_ref): x = pltpu.bitcast(x_ref[...].astype(jnp.float32), jnp.int32) y = pltpu.bitcast(y_ref[...].astype(jnp.float32), jnp.int32) shift = jax.lax.broadcast(16, x.shape) out_ref[...] = pltpu.bitcast( y | jax.lax.shift_right_logical(x, shift), jnp.bfloat16 ) m, n = 16, 128 inp = np.arange(m * n * 2, dtype=jnp.bfloat16).reshape(m, n * 2) x, y = np.split(inp, 2, axis=1) out = self.pallas_call( kernel, out_shape=jax.ShapeDtypeStruct((m * 2, n), jnp.bfloat16), )(x, y) np.testing.assert_array_equal(out, inp.reshape(m * 2, n)) @parameterized.parameters(0, 1) def test_bitcast_vmap(self, axis): def kernel(x_ref, out_ref): out_ref[...] = jax.vmap(functools.partial(pltpu.bitcast, ty=jnp.uint8), in_axes=axis, out_axes=axis)(x_ref[...]) expected = np.random.randint(0, 255, size=(2, 16, 4, 128), dtype=np.uint8) inp = np.ascontiguousarray( expected.swapaxes(-1, -2)).view(np.uint32)[..., 0] assert inp.shape == (2, 16, 128), inp.shape inp = np.moveaxis(inp, 0, axis) expected = np.moveaxis(expected.reshape(2, 16 * 4, 128), 0, axis) out = self.pallas_call( kernel, out_shape=jax.ShapeDtypeStruct(expected.shape, jnp.uint8) )(inp) self.assertArraysEqual(out, expected) @parameterized.parameters([jnp.int32, jnp.int16, jnp.int8, jnp.int4]) def test_row_broadcast(self, dtype): bitwidth = dtypes.itemsize_bits(dtype) if not self.INTERPRET and jtu.get_tpu_version() < 4 and bitwidth < 8: self.skipTest("Requires TPUv4+ for sub-byte types") def kernel(x_ref, y_ref): y_ref[...] = jnp.broadcast_to(x_ref[pl.ds(3, 1)], y_ref.shape).astype(y_ref.dtype) m, n = 4, 1152 x = jax.random.randint( jax.random.key(12), (m, n), minval=-1000, maxval=1000, dtype=jnp.int32 ).astype(dtype) y = self.pallas_call( kernel, out_shape=jax.ShapeDtypeStruct((m, n), jnp.int32) )(x) np.testing.assert_array_equal(y, jnp.broadcast_to(x[3:4], y.shape)) @parameterized.parameters([jnp.uint4, jnp.int4]) def test_tpu_int4_upcast_and_matmul(self, dtype): if not jtu.is_device_tpu_at_least(version=5): self.skipTest("TPUv5+ needed for integer matmuls") def body(x_ref, o_ref): # Test cast from (u)int4 -> int8 ux = lax.convert_element_type(x_ref[...], jnp.int8) o_ref[...] = jax.lax.dot(ux, ux, preferred_element_type=jnp.int32) out = jax.ShapeDtypeStruct((128, 128), jnp.int32) x = jnp.arange(128 * 128, dtype=dtype).reshape((128, 128)) result = self.pallas_call(body, out_shape=out)(x) np.testing.assert_array_equal( result, jax.lax.dot( x.astype(jnp.int8), x.astype(jnp.int8), preferred_element_type=jnp.int32, ), ) def test_sum_of_two_matmuls(self): if not jtu.is_device_tpu_at_least(version=5): self.skipTest("Test requires TPUv5+") M, K = 8, 8 k1, k2, k3, k4 = jax.random.split(jax.random.key(42), 4) a_val = jax.random.normal(k1, (M, K), dtype=jnp.float32) b_val = jax.random.normal(k2, (K,), dtype=jnp.float32) c_val = jax.random.normal(k3, (M, K), dtype=jnp.float32) d_val = jax.random.normal(k4, (K,), dtype=jnp.float32) def kernel(a_ref, b_ref, c_ref, d_ref, o_ref): a = a_ref[:] b = b_ref[:] c = c_ref[:] d = d_ref[:] res1 = jnp.dot(a, b) res2 = jnp.dot(c, d) o_ref[:] = res1 + res2 @jax.jit def pallas_fn(a, b, c, d): return pl.pallas_call( kernel, out_shape=jax.ShapeDtypeStruct((M,), np.float32), grid=(1,), )(a, b, c, d) result_pallas = pallas_fn(a_val, b_val, c_val, d_val) expected = ( jnp.dot(a_val, b_val, precision=jax.lax.Precision.HIGHEST) + jnp.dot(c_val, d_val, precision=jax.lax.Precision.HIGHEST) ) self.assertAllClose(result_pallas, expected, atol=1e-5, rtol=1e-5) @parameterized.product(from_dtype=_JAX_INT_DTYPES, to_dtype=_JAX_INT_DTYPES) def test_integer_cast(self, from_dtype, to_dtype): if ( jnp.iinfo(from_dtype).bits < 8 or jnp.iinfo(to_dtype).bits < 8 ) and not jtu.is_device_tpu_at_least(4): self.skipTest("sub-byte types casting requires TPUv4+") # Generate both low and high values to better cover the entire range # of the source dtype. min_val = from_dtype(jnp.iinfo(from_dtype).min) max_val = from_dtype(jnp.iinfo(from_dtype).max) if jnp.iinfo(from_dtype).bits > 4: x_random = jax.random.randint(jax.random.key(0), shape=(112, 256), minval=min_val, maxval=max_val, dtype=from_dtype) else: # randint does not support sub-byte types. x_random = jnp.arange(112 * 256, dtype=from_dtype).reshape((112, 256)) arange = jnp.arange(8 * 256, dtype=from_dtype).reshape((8, 256)) x = jnp.concatenate([min_val + arange, x_random, max_val - arange], axis=0) def body(x_ref, o_ref): o_ref[...] = lax.convert_element_type(x_ref[...], to_dtype) out = jax.ShapeDtypeStruct(x.shape, to_dtype) expected = x.astype(to_dtype) result = self.pallas_call(body, out_shape=out)(x) np.testing.assert_array_equal(result, expected) def test_select_with_scalar_condition(self): def kernel(cond, lhs, rhs, out): out[:] = jax.lax.select(cond[0] != 0, lhs[:], rhs[:]) def run(cond, lhs, rhs): return self.pallas_call( kernel, out_shape=lhs, grid_spec=pltpu.PrefetchScalarGridSpec( num_scalar_prefetch=0, in_specs=[ pl.BlockSpec(memory_space=pltpu.SMEM), pl.BlockSpec(memory_space=pltpu.VMEM), pl.BlockSpec(memory_space=pltpu.VMEM), ], ), name="select_kernel", )(cond, lhs, rhs) cond = jnp.array([1], dtype=jnp.int32) lhs = jnp.zeros((8, 128), dtype=jnp.float32) rhs = jnp.ones((8, 128), dtype=jnp.float32) assert (run(cond, lhs, rhs) == lhs).all() def test_logical_and_relayouted_mask(self): def get_mask(x_ref): x = x_ref[...] == 1 iota = jax.lax.broadcasted_iota(jnp.int32, x_ref.shape, 1) iota = iota > 7 return jnp.logical_and(x, iota) def body(x_ref, y_ref): y_ref[...] = jnp.where(get_mask(x_ref), 0.0, -1.0) shape = (2, 512) out = jax.ShapeDtypeStruct(shape, jnp.float32) x = jnp.arange(8 * 128, dtype=jnp.int32).reshape(shape) result = self.pallas_call(body, out_shape=out)(x) expected = jnp.ones(x.shape, dtype=jnp.float32) expected = expected.at[...].set(jnp.where(get_mask(x), 0.0, -1.0)) np.testing.assert_array_equal(result, expected) @parameterized.product(dtype=[jnp.float32, jnp.bfloat16, jnp.int16, jnp.int8]) def test_cast_vector_to_mask(self, dtype): if dtype == jnp.int8 and jtu.get_tpu_version() < 5: self.skipTest( f"Not implemented: cast vector to mask with dtype == {dtype}" ) if dtype == jnp.int16: if jtu.get_tpu_version() < 4: self.skipTest("Requires TPUv4+") shape = (128, 128) @functools.partial( self.pallas_call, out_shape=jax.ShapeDtypeStruct(shape, dtype), ) def kernel(x_ref, mask_ref, o_ref): zeros = jnp.zeros_like(x_ref) o_ref[...] = jnp.where(mask_ref[...], x_ref[...], zeros) mask = jax.random.bernoulli(jax.random.key(1234), 0.5, shape).astype(dtype) x = jnp.arange(np.prod(shape), dtype=dtype).reshape(shape) + 1 out = kernel(x, mask) expected = jnp.where(mask, x, jnp.zeros_like(x)) self.assertArraysEqual(out, expected) @parameterized.product( dtype = [jnp.float32, jnp.bfloat16, jnp.int32], axis = [0, 1, 2], reduce_func = [jnp.sum, jnp.max, jnp.min] ) def test_reduction(self, dtype, axis, reduce_func): in_shape = (2, 16, 128) out_shape = list(in_shape) out_shape[axis] = 1 def kernel(x, out): out[:] = reduce_func(x[:], axis, keepdims=True) x = jnp.arange(np.prod(in_shape), dtype=dtype).reshape(in_shape) result = self.pallas_call( kernel, out_shape=jax.ShapeDtypeStruct(out_shape, x.dtype), )(x) expected = reduce_func(x, axis, keepdims=True) np.testing.assert_array_equal(result, expected) @parameterized.product( axis=[0, 1, 2], in_shape=[(2, 29, 206), (12, 28), (12,)], reduce_func = [jnp.argmax, jnp.argmin], keepdims=[False, True], ) def test_reduce_index(self, axis, in_shape, reduce_func, keepdims): dtype = jnp.float32 rank = len(in_shape) if axis >= rank: self.skipTest("Requires axis < rank") if axis == rank - 1: if keepdims and not jtu.is_device_tpu_at_least(version=4): self.skipTest("Requires TPUv4+ for axis=rank-1 and keepdims=True") if not keepdims and not jtu.is_device_tpu_at_least(version=5): self.skipTest("Requires TPUv5+ for axis=rank-1 and keepdims=False") if rank == 1 and not keepdims: self.skipTest("Scalar output not supported") out_shape = list(in_shape) if keepdims: out_shape[axis] = 1 else: del out_shape[axis] def kernel(x, out): out[:] = reduce_func(x[:], axis, keepdims=keepdims) x = jax.random.permutation( jax.random.key(22), jnp.arange(np.prod(in_shape), dtype=dtype) ).reshape(in_shape) result = self.pallas_call( kernel, out_shape=jax.ShapeDtypeStruct(out_shape, jnp.int32), )(x) expected = reduce_func(x, axis, keepdims=keepdims) np.testing.assert_array_equal(result, expected) @parameterized.product( in_shape=[(2, 29, 206), (12, 28), (32, 512)], axis=[0, 1, 2], reduce_func=[ jnp.argmax, jnp.argmin, functools.partial(lax.argmax, index_dtype=jnp.int16), functools.partial(lax.argmin, index_dtype=jnp.int16), ], ) def test_bf16_input_reduce_index(self, in_shape, axis, reduce_func): if not jtu.is_device_tpu_at_least(6): self.skipTest("Requires TPUv6+") if axis >= len(in_shape): self.skipTest("Requires axis < rank") dtype = jnp.bfloat16 if reduce_func in (jnp.argmax, jnp.argmin): index_dtype = jnp.int32 else: index_dtype = jnp.int16 out_shape = list(in_shape) del out_shape[axis] def kernel(x, out): out[:] = reduce_func(x[:], axis) x = jnp.argsort( jax.random.normal(jax.random.key(22), shape=in_shape), axis=axis, ).astype(dtype) result = self.pallas_call( kernel, out_shape=jax.ShapeDtypeStruct(out_shape, dtype=index_dtype), )(x) np.testing.assert_array_equal(result, reduce_func(x, axis)) @parameterized.product( k=[1, 2, 4, 8], in_shape=[(13, 32, 512), (10, 16, 128)], axis=[0, 1, 2, -1, -2, -3], dtype=[jnp.float32, jnp.bfloat16], ) def test_top_k(self, k, in_shape, axis, dtype): if not jtu.is_device_tpu_at_least(4): self.skipTest("Requires TPUv4+") if dtype == jnp.bfloat16 and not jtu.is_device_tpu_at_least(6): self.skipTest("Requires TPUv6+ for bfloat16") out_val_shape = list(in_shape) out_val_shape[axis] = k def kernel(x_ref, val_ref, idx_ref): val_ref[...], idx_ref[...] = jax.lax.top_k( x_ref[...], k=k, axis=axis, is_stable=False ) # Ensure unique values in x so that no ties need to be broken. idx = jnp.argsort( jax.random.normal(jax.random.key(42), shape=in_shape), axis=axis ).astype(jnp.float32) x = (jnp.exp2(idx // 128.0) * (1.0 + (idx % 128.0) / 128.0)).astype(dtype) res_v, res_i = self.pallas_call( kernel, out_shape=[ jax.ShapeDtypeStruct(out_val_shape, dtype), jax.ShapeDtypeStruct(out_val_shape, jnp.int32), ], )(x) exp_v, exp_i = jax.lax.top_k(x, k=k, axis=axis) np.testing.assert_array_equal(res_v, exp_v) np.testing.assert_array_equal(res_i, exp_i) @parameterized.product( k=[1, 2, 4, 8], in_shape=[(13, 32, 512), (10, 16, 1024), (9, 11, 300)], axis=[0, 1, 2, -1, -2, -3], dtype=[jnp.float32, jnp.bfloat16], recall_target=[0.95, 0.99, 1.0], ) def test_approx_max_k(self, k, in_shape, axis, dtype, recall_target): if not jtu.is_device_tpu_at_least(4): self.skipTest("Requires TPUv4+") if dtype == jnp.bfloat16 and not jtu.is_device_tpu_at_least(6): self.skipTest("Requires TPUv6+ for bfloat16") out_val_shape = list(in_shape) out_val_shape[axis] = k def kernel(x_ref, val_ref, idx_ref): val_ref[...], idx_ref[...] = jax.lax.approx_max_k( x_ref[...], k=k, reduction_dimension=axis, recall_target=recall_target ) # Ensure unique values in x so that no ties need to be broken. idx = jnp.argsort( jax.random.normal(jax.random.key(42), shape=in_shape), axis=axis ).astype(jnp.float32) x = (jnp.exp2(idx // 128.0) * (1.0 + (idx % 128.0) / 128.0)).astype(dtype) res_v, res_i = self.pallas_call( kernel, out_shape=[ jax.ShapeDtypeStruct(out_val_shape, dtype), jax.ShapeDtypeStruct(out_val_shape, jnp.int32), ], )(x) # Gather consistency: returned values must match elements at returned indices gathered_v = np.take_along_axis(np.asarray(x), np.asarray(res_i), axis=axis) np.testing.assert_array_equal(res_v, gathered_v) # Descending sorted order sorted_v = np.flip(np.sort(np.asarray(res_v), axis=axis), axis=axis) np.testing.assert_array_equal(res_v, sorted_v) # Statistical Recall check vs ground-truth exact top-k _, gt_i = jax.lax.top_k(x, k=k, axis=axis) res_i_flat = np.moveaxis(np.asarray(res_i), axis, -1).reshape(-1, k) gt_i_flat = np.moveaxis(np.asarray(gt_i), axis, -1).reshape(-1, k) hits = sum(len(set(r) & set(g)) for r, g in zip(res_i_flat, gt_i_flat)) recall = hits / gt_i_flat.size self.assertGreaterEqual(recall, recall_target) @parameterized.product( shape=[(129, 129), (1, 129), (2, 129), (4, 129)], msk_dtype=[jnp.float32, jnp.bfloat16, jnp.int8], dtype=[jnp.float32, jnp.bfloat16], ) def test_i1_relayout_bw(self, shape, msk_dtype, dtype): msk_bitwidth = dtypes.itemsize_bits(msk_dtype) if jtu.get_tpu_version() < 5 and msk_bitwidth < 16: self.skipTest( "Not implemented: cast vector to mask with bitwidth ==" f" {msk_bitwidth}" ) @functools.partial( pl.pallas_call, out_shape=jax.ShapeDtypeStruct(shape, dtype), ) def kernel(x_ref, mask_ref, o_ref): zeros = jnp.zeros_like(x_ref) o_ref[...] = jnp.where(mask_ref[...], x_ref[...], zeros) mask = jax.random.bernoulli(jax.random.key(1234), 0.5, shape).astype( msk_dtype ) x = jnp.arange(np.prod(shape), dtype=dtype).reshape(shape) + 1 out = kernel(x, mask) expected = jnp.where(mask, x, jnp.zeros_like(x)) self.assertArraysEqual(out, expected) @parameterized.product( msk_dtype=[jnp.float32, jnp.bfloat16, jnp.int8], dtype=[jnp.float32, jnp.bfloat16, jnp.int8], ) def test_i1_relayout_bw_tiling(self, msk_dtype, dtype): shape = (256, 256) msk_bitwidth = dtypes.itemsize_bits(msk_dtype) msk_packing = 32 // msk_bitwidth if jtu.get_tpu_version() < 5 and msk_bitwidth < 16: self.skipTest( "Not implemented: cast vector to mask with bitwidth ==" f" {msk_bitwidth}" ) # Creating large tiling for masks by passing i32 vector first and # then bitcast to msk_dtype so the tiling is also bitcasted from # T(8, 128) to T(8 * msk_packing, 128). @functools.partial( pl.pallas_call, out_shape=jax.ShapeDtypeStruct(shape, dtype), ) def kernel(x_ref, msk_ref, o_ref): zeros = jnp.zeros_like(x_ref) msk = pltpu.bitcast(msk_ref[...], msk_dtype) o_ref[...] = jnp.where(msk, x_ref[...], zeros) mask = jax.random.bernoulli(jax.random.key(1234), 0.5, shape).astype( msk_dtype ) if msk_bitwidth < 32: mask_for_bitcast = mask.reshape( shape[0] // msk_packing, msk_packing, shape[1] ).swapaxes(-1, -2) else: mask_for_bitcast = mask mask_i32 = lax.bitcast_convert_type( mask_for_bitcast, jnp.int32, ) x = jnp.arange(np.prod(shape), dtype=dtype).reshape(shape) + 1 out = kernel(x, mask_i32) expected = jnp.where(mask, x, jnp.zeros_like(x)) self.assertArraysEqual(out, expected) @parameterized.product( msk_dtype=[jnp.int32, jnp.int16, jnp.int8], dtype=[jnp.int32, jnp.int16, jnp.int8], ) def test_i1_relayout_bw_1d_tiling(self, msk_dtype, dtype): if msk_dtype == jnp.int8 and jtu.get_tpu_version() < 5: self.skipTest("Requires TPUv5+") if msk_dtype == jnp.int16: if jtu.get_tpu_version() < 4: self.skipTest("Requires TPUv4+") shape = (1024,) @functools.partial( pl.pallas_call, out_shape=jax.ShapeDtypeStruct(shape, dtype), ) def kernel(x_ref, mask_ref, o_ref): zeros = jnp.zeros_like(x_ref) o_ref[...] = jnp.where(mask_ref[...], x_ref[...], zeros) mask = jax.random.bernoulli(jax.random.key(1234), 0.5, shape).astype( msk_dtype ) x = jnp.arange(np.prod(shape), dtype=dtype).reshape(shape) + 1 out = kernel(x, mask) expected = jnp.where(mask, x, jnp.zeros_like(x)) self.assertArraysEqual(out, expected) @parameterized.product( target=(jnp.int8,), # TODO(apaszke): Add int4. round=(False, True), ) def test_quantize(self, target, round): if not jtu.is_device_tpu_at_least(version=6): self.skipTest("Requires TPUv6+") shape = (256, 256) # NOTE: 256 * 256 == 2 ** 16, so those are all bf16 values. x = lax.bitcast_convert_type( np.arange(math.prod(shape), dtype=jnp.uint16).reshape(shape), jnp.bfloat16, ) round_fn = jnp.rint if round else lambda x: x def kernel(x_ref, o_ref): o_ref[...] = round_fn(x_ref[...]).astype(target) out = self.pallas_call( kernel, out_shape=jax.ShapeDtypeStruct(shape, target) )(x) ref = jax.jit(lambda x: round_fn(x).astype(target))(x) np.testing.assert_array_equal(out, ref) @parameterized.product( idx_shape=[(8, 128), (2, 32, 256)], axis=[-2, -1], dtype=[jnp.int32, jnp.int16, jnp.int8], mode=["promise_in_bounds", None], wrap_negative_indices=[True, False], ) def test_dynamic_gather_along_axis( self, idx_shape, axis, dtype, mode, wrap_negative_indices ): if (axis == -2 and not jtu.is_device_tpu_at_least(version=5)) or ( axis == -1 and not jtu.is_device_tpu_at_least(version=4) ): self.skipTest("Requires TPUv5+ for axis=-2 and TPUv4+ for axis=-1") if dtype != jnp.int32: if wrap_negative_indices: self.skipTest( "With wrap_negative_indices=True, jnp.take_along_axis will cast" " indices to int32 causing a type mismatch for dynamic gather" ) if ( dtype == jnp.int16 or (dtype == jnp.int8 and axis == -2) ) and not jtu.is_device_tpu_at_least(version=6): self.skipTest("Requires TPUv6+") x_shape = list(idx_shape) # Only a single vreg along the gather axis is supported bitwidth = dtypes.itemsize_bits(dtype) x_shape[axis] = (8 * 32 // bitwidth, 128)[axis] def kernel(x, indices, out): out[...] = jnp.take_along_axis( x[...], indices[...], axis, mode=mode, wrap_negative_indices=wrap_negative_indices, ) x = np.arange(np.prod(x_shape), dtype=dtype).reshape(x_shape) idx = jax.random.randint( key=jax.random.key(1234), shape=idx_shape, minval=0, maxval=x_shape[axis], dtype=dtype, ) actual = self.pallas_call( kernel, out_shape=jax.ShapeDtypeStruct(idx_shape, dtype) )(x, idx) expected = np.take_along_axis(x, idx, axis=axis) np.testing.assert_array_equal(actual, expected) @parameterized.product(dtype=[jnp.float32, jnp.bfloat16]) def test_float_div(self, dtype): kwargs = {} if jtu.is_device_tpu_at_least(version=6): kwargs.update(dict(rtol=1e-2)) def kernel(x, y, out): out[:] = jax.lax.div(x[:], y[:]) run = pl.pallas_call( kernel, out_shape=jax.ShapeDtypeStruct((8, 128), dtype), ) k1, k2 = jax.random.split(jax.random.key(1234), 2) x = jax.random.normal(k1, (8, 128), dtype=dtype) y = jax.random.normal(k2, (8, 128), dtype=dtype) np.testing.assert_allclose(run(x, y), jax.lax.div(x, y), **kwargs) @parameterized.product( dtype=[jnp.float32, jnp.bfloat16, jnp.int8], ) def test_concat_mask(self, dtype): bitwidth = dtypes.itemsize_bits(dtype) if jtu.get_tpu_version() < 5 and bitwidth < 16: self.skipTest( f"Not implemented: cast vector to mask with bitwidth == {bitwidth}" ) shape = (128, 128) def kernel(x, out): mask = x[...] != 0 concated_mask = jnp.concatenate([mask, mask], axis=0) concated_x = jnp.concatenate([x[:], x[:]], axis=0) out[:] = lax.select(concated_mask, concated_x, jnp.zeros_like(concated_x)) x = jax.random.normal(jax.random.key(1234), shape, dtype=jnp.float32) if dtype == jnp.int8: x = (x * 100).astype(jnp.int8) else: x = x.astype(dtype) out = self.pallas_call( kernel, out_shape=jax.ShapeDtypeStruct((shape[0] * 2, shape[1]), dtype) )(x) concated_mask = jnp.concatenate([x != 0, x != 0], axis=0) concated_x = jnp.concatenate([x, x], axis=0) expected = lax.select(concated_mask, concated_x, jnp.zeros_like(concated_x)) np.testing.assert_array_equal(out, expected) def test_reduce_with_const(self): m = 1 d = 1024 x = jnp.ones((m, d), jnp.bfloat16) def dot(x, y): return jax.lax.dot_general( x, y, (((1,), (1,)), ((), ())), preferred_element_type=jnp.float32, ) def kernel(x, out): out[:] = dot(x[:], jnp.ones((1, d), jnp.bfloat16)) run = pl.pallas_call(kernel, jax.ShapeDtypeStruct((m, 1), jnp.float32)) output = run(x) expected = dot(x[:], jnp.ones((1, d), jnp.bfloat16)) np.testing.assert_array_equal(output, expected) # We need to manually run the test with the env variable # `export LIBTPU_INIT_ARGS="--xla_jf_bounds_check=true"` def test_disable_bounds_check(self): if jtu.get_tpu_version() < 4: self.skipTest("Requires TPUv4+") src_shape = (8, 128) tgt_shape = (16, 256) def kernel(src, tgt): tgt[:] = src[ tuple( pl.ds(pl.program_id(0) * s, d) for s, d in zip(src.shape, tgt.shape) ) ] x = jnp.arange(np.prod(src_shape), dtype=jnp.float32).reshape(src_shape) run = pl.pallas_call( kernel, jax.ShapeDtypeStruct(tgt_shape, jnp.float32), grid=(1,), compiler_params=pltpu.CompilerParams(disable_bounds_checks=True), ) output = run(x) np.testing.assert_array_equal( output[tuple(slice(0, d) for d in src_shape)], x ) def test_while_loop_arg_num_change(self): # This kernel will generate a while loop that will be CSEd by MLIR to have # the different number of argments in before region and after region. def kernel( out_ref, a, ): def loop_cond(state): _, y = state return y def loop_body(state): x, y = state def then_0(): def then_1(): return jnp.int32(0) def else_1(): a[0] = a[0] + 1 return jnp.int32(1) z = lax.cond(x == 0, then_1, else_1) new_x = z new_y = z != 0 return new_x, new_y def else_0(): return x, jnp.bool_(False) new_x, new_y = lax.cond(y, then_0, else_0) return (new_x, new_y) out_ref[0] = lax.while_loop( loop_cond, loop_body, (jnp.int32(0), jnp.bool_(True)) )[0] output = pl.pallas_call( kernel, out_shape=jax.ShapeDtypeStruct((1,), jnp.int32), in_specs=(), out_specs=pl.BlockSpec(memory_space=pltpu.SMEM), scratch_shapes=(pltpu.SMEM((1,), jnp.int32),), )()[0] self.assertEqual(output, 0) def test_produce_predicate_phi(self): def kernel( out_ref, a, ): def loop_cond(state): x, y = state return jnp.logical_or(y, (x == 1)) def loop_body(state): x, y = state def then_0(): def then_1(): return jnp.int32(0) def else_1(): a[0] = a[0] + 1 return jnp.int32(1) z = lax.cond(x == 0, then_1, else_1) new_x = z new_y = z != 0 return new_x, new_y def else_0(): return x, jnp.bool_(False) new_x, new_y = lax.cond(y, then_0, else_0) return (new_x, new_y) out_ref[0] = lax.while_loop( loop_cond, loop_body, (jnp.int32(0), jnp.bool_(True)) )[0] output = pl.pallas_call( kernel, out_shape=jax.ShapeDtypeStruct((1,), jnp.int32), in_specs=(), out_specs=pl.BlockSpec(memory_space=pltpu.SMEM), scratch_shapes=(pltpu.SMEM((1,), jnp.int32),), )()[0] self.assertEqual(output, 0) def test_retiling_with_replicated_lane(self): shape = (32, 1) broadcast_shape = (32, 256) @functools.partial( pl.pallas_call, out_shape=jax.ShapeDtypeStruct((8, 4, 256), jnp.float32), ) def kernel(x_ref, o_ref): o_ref[...] = jnp.broadcast_to( x_ref[...], broadcast_shape ).reshape(o_ref.shape) x = jnp.arange(np.prod(shape), dtype=jnp.float32).reshape(shape) out = kernel(x).reshape(broadcast_shape) expected = jnp.broadcast_to(x, broadcast_shape) np.testing.assert_array_equal(out, expected) @parameterized.parameters( [jnp.bfloat16, jnp.float8_e5m2, jnp.float8_e4m3fn, jnp.float8_e4m3b11fnuz] ) def test_stochastic_round(self, target_dtype): if not jtu.is_device_tpu_at_least(version=5): self.skipTest("Requires TPU v5+") def kernel(x_ref, b_ref, o_ref): o_ref[...] = pltpu.stochastic_round( x_ref[...], b_ref[...], target_dtype=target_dtype ) shape = (8, 128) k1, k2 = jax.random.split(jax.random.key(4242), 2) x = jax.random.normal(k1, shape, dtype=jnp.float32) bits = jax.random.bits(k2, shape, dtype=jnp.uint32) x_cast = x.astype(target_dtype) x_cast_as_f32 = x_cast.astype(jnp.float32) max_val = jnp.finfo(target_dtype).max min_val = jnp.finfo(target_dtype).min lower = jnp.where(x_cast_as_f32 > x, jnp.nextafter(x_cast, min_val), x_cast) upper = jnp.where(x_cast_as_f32 < x, jnp.nextafter(x_cast, max_val), x_cast) result = pl.pallas_call( kernel, out_shape=jax.ShapeDtypeStruct(x.shape, target_dtype), )(x, bits) int_dtype = getattr(jnp, f"uint{dtypes.itemsize_bits(target_dtype)}") is_correct_bitwise = ( (result.view(int_dtype) == lower.view(int_dtype)) | (result.view(int_dtype) == upper.view(int_dtype)) ) is_correct = jnp.where( jnp.isnan(x_cast), jnp.isnan(result), is_correct_bitwise ) self.assertTrue(jnp.all(is_correct)) def _pack_unpack_elementwise_test_data( self, shape, unpacked_dtype, packed_dtype): """Generates data for test_pack_elementwise and test_unpack_elementwise.""" unpacked_bitwidth = dtypes.itemsize_bits(unpacked_dtype) packed_bitwidth = dtypes.itemsize_bits(packed_dtype) num_sources = unpacked_bitwidth // packed_bitwidth stacked_sources = jnp.asarray(rand((num_sources, *shape), unpacked_dtype)) stacked_results = ( stacked_sources.astype(packed_dtype) .view(getattr(jnp, f"uint{packed_bitwidth}")) .astype(getattr(jnp, f"uint{unpacked_bitwidth}")) ) shifts = jnp.arange(num_sources, dtype=jnp.uint32) * packed_bitwidth shifts = jnp.expand_dims(shifts, axis=tuple(range(1, stacked_results.ndim))) packed_data = jnp.bitwise_or.reduce( stacked_results.astype(jnp.uint32) << shifts, axis=0 ).astype(getattr(jnp, f"uint{unpacked_bitwidth}")) return stacked_sources, packed_data @parameterized.product( config=[ (jnp.float32, jnp.bfloat16), (jnp.int32, jnp.int16), (jnp.int32, jnp.int8), (jnp.int32, jnp.int4), (jnp.int16, jnp.int8), (jnp.int8, jnp.int4), (jnp.int4, jnp.int2), ], shape=[(8, 128), (2, 15, 300), (512, 256)], ) def test_pack_elementwise(self, config, shape): unpacked_dtype, packed_dtype = config if packed_dtype == jnp.int2: if not jtu.is_device_tpu_at_least(version=5): self.skipTest("Requires TPU v5+") if (shape[-2] % (8 * 16)) or (shape[-1] % 128): raise self.skipTest( "int2 is only supported for shapes with vreg alignment" ) src_bitwidth = dtypes.itemsize_bits(unpacked_dtype) tgt_bitwidth = dtypes.itemsize_bits(packed_dtype) num_sources = src_bitwidth // tgt_bitwidth output_dtype = getattr(jnp, f"uint{src_bitwidth}") def kernel(xs_ref, o_ref): xs = [xs_ref[i] for i in range(num_sources)] o_ref[...] = pltpu.pack_elementwise(xs, packed_dtype=packed_dtype) stacked_sources, expected = self._pack_unpack_elementwise_test_data( shape, unpacked_dtype, packed_dtype ) result = self.pallas_call( kernel, out_shape=jax.ShapeDtypeStruct(shape, output_dtype), )(stacked_sources) np.testing.assert_array_equal(result, expected) @parameterized.product( config=[ (jnp.float32, jnp.bfloat16), (jnp.int32, jnp.int16), (jnp.int32, jnp.int8), (jnp.int32, jnp.int4), (jnp.int16, jnp.int8), (jnp.int16, jnp.int4), (jnp.int8, jnp.int4), (jnp.int4, jnp.int2), ], index=[0, 1, 2, 3], shape=[(8, 128), (2, 15, 300), (512, 256)], ) def test_unpack_elementwise(self, config, index, shape): unpacked_dtype, packed_dtype = config if packed_dtype == jnp.bfloat16 and not jtu.is_device_tpu_at_least( version=5 ): self.skipTest("Requires TPU v5+") if packed_dtype == jnp.int2: if not jtu.is_device_tpu_at_least(version=5): self.skipTest("Requires TPU v5+") if (shape[-2] % (8 * 16)) or (shape[-1] % 128): raise self.skipTest( "int2 is only supported for shapes with vreg alignment" ) unpacked_bitwidth = dtypes.itemsize_bits(unpacked_dtype) packed_bitwidth = dtypes.itemsize_bits(packed_dtype) packing_factor = unpacked_bitwidth // packed_bitwidth if index >= packing_factor: self.skipTest( f"Index {index} out of bounds for packing factor {packing_factor}") def kernel(x_ref, o_ref): o_ref[...] = pltpu.unpack_elementwise( x_ref[...], index=index, packed_dtype=packed_dtype, unpacked_dtype=unpacked_dtype ) sources, packed = self._pack_unpack_elementwise_test_data( shape, unpacked_dtype, packed_dtype ) expected = sources[index].astype(packed_dtype).astype(unpacked_dtype) result = self.pallas_call( kernel, out_shape=jax.ShapeDtypeStruct(shape, unpacked_dtype), )(packed) np.testing.assert_array_equal(result, expected) @parameterized.product( dtype=[jnp.float32, jnp.bfloat16, jnp.int32, jnp.int8], shape_reps=[ ((8, 128), (2, 1)), ((8, 128), (1, 2)), ((8, 128), (2, 2)), ((128, 8, 128), (3,2,1)), ], ) def test_tile(self, dtype, shape_reps): shape, reps = shape_reps k1 = jax.random.key(1234) if jnp.issubdtype(dtype, jnp.integer): x = jax.random.randint( k1, shape, minval=dtypes.iinfo(dtype).min, maxval=dtypes.iinfo(dtype).max, dtype=dtype, ) else: x = jax.random.normal(k1, shape, dtype=dtype) def kernel(x_ref, y_ref): y_ref[...] = jnp.tile(x_ref[...], reps) run = pl.pallas_call( kernel, out_shape=jax.ShapeDtypeStruct(np.multiply(shape, reps), dtype), ) output = run(x) expected = jnp.tile(x, reps) np.testing.assert_array_equal(output, expected) @parameterized.parameters( jnp.float32, jnp.bfloat16, ) def test_sigmoid(self, dtype): shape = (32, 128) x = jax.random.normal(jax.random.key(42), shape, dtype=dtype) @functools.partial( self.pallas_call, out_shape=jax.ShapeDtypeStruct(shape, dtype), ) def kernel(x_ref, o_ref): o_ref[...] = jax.nn.sigmoid(x_ref[...]) self.assertAllClose( kernel(x), jax.nn.sigmoid(x), ) def test_erf(self): # Regression test for https://github.com/jax-ml/jax/issues/36149. x = jnp.zeros((32,), dtype=jnp.float32) @functools.partial( self.pallas_call, out_shape=jax.ShapeDtypeStruct.like(x), ) def kernel(x_ref, out_ref): out_ref[...] = jax.scipy.special.erf(x_ref[...]) _ = jax.jit(kernel).lower(x) @parameterized.parameters(jnp.uint4, jnp.uint8, jnp.uint16, jnp.uint32) def test_unsigned_dtype_dot_raises(self, dtype): k = 256 packing = 32 // jnp.iinfo(dtype).bits lhs = jnp.zeros((8 * packing, k), dtype=dtype) rhs = jnp.zeros((k, 128), dtype=dtype) def kernel(lhs_ref, rhs_ref, o_ref): o_ref[...] = lhs_ref[...] @ rhs_ref[...] out_shape = jax.ShapeDtypeStruct((8 * packing, 128), dtype) with self.assertRaisesRegex( NotImplementedError, "Unsigned integer dtype.*dot_general.*matmul" ): self.pallas_call(kernel, out_shape=out_shape)(lhs, rhs) @parameterized.product( dtype=[jnp.int32, jnp.int16, jnp.int8], shapes_axis_and_out=[ ((11, 300), (9, 300), 0, (20, 300)), ((20, 400), (17, 400), 0, (37, 400)), ((35, 200), (41, 200), 0, (76, 200)), ((35, 200), (35, 300), 1, (35, 500)), ], ) def test_2d_concatenate(self, dtype, shapes_axis_and_out): packing = 32 // jnp.iinfo(dtype).bits if packing > 1 and not jtu.is_device_tpu_at_least(version=4): self.skipTest("Requires TPU v4+") if packing > 2 and not jtu.is_device_tpu_at_least(version=5): self.skipTest("Requires TPU v5+") x_shape, y_shape, axis, out_shape = shapes_axis_and_out min_val = dtype(jnp.iinfo(dtype).min) max_val = dtype(jnp.iinfo(dtype).max) x = jax.random.randint( jax.random.key(42), x_shape, minval=min_val, maxval=max_val, dtype=dtype ) y = jax.random.randint( jax.random.key(0), y_shape, minval=min_val, maxval=max_val, dtype=dtype ) @functools.partial( self.pallas_call, out_shape=jax.ShapeDtypeStruct(out_shape, dtype), ) def kernel(x_ref, y_ref, o_ref): o_ref[...] = jnp.concatenate([x_ref[...], y_ref[...]], axis=axis) np.testing.assert_array_equal( kernel(x, y), jnp.concatenate([x, y], axis=axis), ) @parameterized.product( dtype=[jnp.int32, jnp.int16, jnp.int8], shapes_and_out=[ # cases for 32-bit (200, 300, 500), (1153, 1200, 2353), # cases for 16-bit (256, 1024, 1280), (2176, 2100, 4276), (2304, 2560, 4864), # cases for 8-bit (512, 2048, 2560), (4225, 5000, 9225), (4608, 5120, 9728), ], ) def test_1d_concatenate(self, dtype, shapes_and_out): packing = 32 // jnp.iinfo(dtype).bits x_shape, y_shape, out_shape = shapes_and_out if packing >= 4: if not jtu.is_device_tpu_at_least(version=5): self.skipTest("Requires TPU v5+") if packing >= 2: if not jtu.is_device_tpu_at_least(version=4): self.skipTest("Requires TPU v4+") min_val = dtype(jnp.iinfo(dtype).min) max_val = dtype(jnp.iinfo(dtype).max) x = jax.random.randint( jax.random.key(42), (x_shape,), minval=min_val, maxval=max_val, dtype=dtype, ) y = jax.random.randint( jax.random.key(0), (y_shape,), minval=min_val, maxval=max_val, dtype=dtype, ) @functools.partial( self.pallas_call, out_shape=jax.ShapeDtypeStruct((out_shape,), dtype), ) def kernel(x_ref, y_ref, o_ref): o_ref[...] = jnp.concatenate([x_ref[...], y_ref[...]], axis=0) np.testing.assert_array_equal( kernel(x, y), jnp.concatenate([x, y], axis=0), ) @parameterized.parameters(jnp.int16, jnp.uint16, jnp.bfloat16) def test_concatenate_16bit_slices(self, dtype): if jtu.is_device_tpu_at_least(version=5) and not jtu.is_libtpu_at_least( "0.0.50" ): self.skipTest("Requires libtpu >= 0.0.50") if dtype == jnp.bfloat16 and not jtu.is_device_tpu_at_least(version=4): self.skipTest( "v4- does not have native vpackc support, so subnormals will be" " flushed to zero." ) b, r, l = np.ogrid[0:2, 0:8, 0:128] x_i16 = (1000 * (r - 2) + 100 * b + l).astype(np.int16) x_np = x_i16.view(dtype) x = jnp.asarray(x_np) @functools.partial( self.pallas_call, out_shape=jax.ShapeDtypeStruct((2, 2, 100), dtype), ) def kernel(x_ref, o_ref): x_val = x_ref[...] o_ref[...] = jnp.concatenate( [x_val[:, 0:1, 0:100], x_val[:, 2:3, 20:120]], axis=1 ) out = np.asarray(kernel(x)) expected = np.concatenate( [x_np[:, 0:1, 0:100], x_np[:, 2:3, 20:120]], axis=1 ) np.testing.assert_array_equal(out.view(np.uint16), expected.view(np.uint16)) def test_fuse_transposed_lhs_in_matmul(self): lhs_shape = (512, 128) rhs_shape = (512, 256) def kernel(lhs_ref, rhs_ref, o_ref): o_ref[...] = jnp.einsum("km,kn->mn", lhs_ref[...], rhs_ref[...]) lhs_key, rhs_key = jax.random.split(jax.random.key(42), 2) lhs = ( jax.random.normal(lhs_key, lhs_shape, dtype=jnp.float32) .astype(jnp.bfloat16) .astype(jnp.float32) ) rhs = ( jax.random.normal(rhs_key, rhs_shape, dtype=jnp.float32) .astype(jnp.bfloat16) .astype(jnp.float32) ) result = self.pallas_call( kernel, out_shape=jax.ShapeDtypeStruct( (lhs_shape[1], rhs_shape[1]), jnp.float32 ), compiler_params=pltpu.CompilerParams( fuse_transposed_lhs_in_matmul=True ), )(lhs, rhs) np.testing.assert_allclose( result, jnp.einsum("km,kn->mn", lhs, rhs), atol=1e-5 ) @parameterized.product( dtype=[jnp.int8, jnp.int16, jnp.uint8, jnp.uint16], op=[lax.eq, lax.ne, lax.lt, lax.gt, lax.le, lax.ge], ) def test_scalar_comparison(self, dtype, op): @functools.partial( self.pallas_call, out_shape=jax.ShapeDtypeStruct((1,), jnp.bool_), in_specs=[ pl.BlockSpec(memory_space=pltpu.SMEM), pl.BlockSpec(memory_space=pltpu.SMEM), ], out_specs=pl.BlockSpec(memory_space=pltpu.SMEM), ) def _kernel(x_ref, y_ref, o_ref): o_ref[0] = op(x_ref[0].astype(jnp.int32), y_ref[0].astype(jnp.int32)) if jnp.issubdtype(dtype, jnp.signedinteger): x_val = -1 y_val = 0 else: x_val = 255 if dtype == jnp.uint8 else 65535 y_val = 0 x = jnp.array([x_val], dtype=dtype) y = jnp.array([y_val], dtype=dtype) out = _kernel(x, y) expected = op(x, y) self.assertAllClose(out, expected) @parameterized.product(approx=[False, True], full_range=[False, True], dtype=[jnp.bfloat16, jnp.float32]) def test_reciprocal(self, approx, full_range, dtype): if dtype == jnp.bfloat16: if not jtu.is_device_tpu_at_least(6): self.skipTest("BF16 reciprocal not supported before v6e.") if approx: self.skipTest("BF16 reciprocal does not support approx=True.") if not full_range: self.skipTest("BF16 reciprocal does not support full_range=False.") shape = (32, 256) x = jnp.arange(np.prod(shape), dtype=dtype).reshape(shape) if not full_range: x = jnp.where(x == 0, -1.0, x) def kernel(x_ref, o_ref): o_ref[...] = pltpu.reciprocal( x_ref[...], approx=approx, full_range=full_range ) out = self.pallas_call( kernel, out_shape=jax.ShapeDtypeStruct(shape, dtype) )(x) kwargs = {} if approx: kwargs.update(dict(atol=2e-5, rtol=2e-5)) np.testing.assert_allclose(out, jax.lax.reciprocal(x), **kwargs) @parameterized.product( bits=[ (2,), (4,), (4, 2), ] ) def test_join_bits(self, bits): if 2 in bits and not jtu.is_device_tpu_at_least(version=5): self.skipTest("Requires TPU v5+") if 4 in bits and not jtu.is_device_tpu_at_least(version=4): self.skipTest("Requires TPU v4+") shape = (256, 128) i8_inputs = jax.random.randint( jax.random.key(0), shape, 0, 2**8, dtype=jnp.uint8, ) bit_inputs = [] shift = 0 for bit in bits: mask = (1 << bit) - 1 dtype = getattr(jnp, f"int{bit}") bit_inputs.append(((i8_inputs >> shift) & mask).astype(dtype)) shift += bit kw_pairs = [(4, "four_bits"), (2, "two_bits")] def kernel(*args): *bit_args, o_ref = args kwargs = {} bit_args_list = list(bit_args) for bit, kw in kw_pairs: if bit in bits: kwargs[kw] = bit_args_list.pop(0)[...] o_ref[...] = pltpu.join_bits(**kwargs) result = self.pallas_call( kernel, out_shape=jax.ShapeDtypeStruct(shape, jnp.int8), )(*bit_inputs) total_bits = sum(bits) mask = (1 << total_bits) - 1 np.testing.assert_array_equal(result.astype(jnp.uint8) >> total_bits, 0) np.testing.assert_array_equal(result & mask, i8_inputs & mask) @parameterized.product( bitwidths=[ (4, 2), (2, 4), (4,), (2,), ] ) def test_split_bits(self, bitwidths): if 2 in bitwidths and not jtu.is_device_tpu_at_least(version=5): self.skipTest("Requires TPU v5+") if 4 in bitwidths and not jtu.is_device_tpu_at_least(version=4): self.skipTest("Requires TPU v4+") shape = (256, 128) i8_inputs = jax.random.randint( jax.random.key(0), shape, 0, 2**8, dtype=jnp.uint8, ) def kernel(x_ref, *o_refs): outs = pltpu.split_bits(x_ref[...], bitwidths=bitwidths) for o_ref, out in zip(o_refs, outs): o_ref[...] = out out_shapes = [ jax.ShapeDtypeStruct(shape, getattr(jnp, f"int{bw}")) for bw in bitwidths ] if len(out_shapes) == 1: result = self.pallas_call( kernel, out_shape=out_shapes[0], )(i8_inputs) results = [result] else: results = self.pallas_call( kernel, out_shape=tuple(out_shapes), )(i8_inputs) expected = [] shift = 0 for bw in bitwidths: mask = (1 << bw) - 1 dtype = getattr(jnp, f"int{bw}") expected.append(((i8_inputs >> shift) & mask).astype(dtype)) shift += bw for res, exp in zip(results, expected): np.testing.assert_array_equal(res, exp) def test_fp6_unpack(self): if not jtu.is_device_tpu_at_least(version=5): self.skipTest("Requires TPU v5+") shape = (256, 128) two_bits_in = jax.random.randint( jax.random.key(0), shape, -2, 2, dtype=jnp.int2 ) four_bits_in = jax.random.randint( jax.random.key(1), shape, -8, 8, dtype=jnp.int4 ) def unpack_kernel(two_b_ref, four_b_ref, o_ref): o_ref[...] = pltpu.unpack_fp6(two_b_ref[...], four_b_ref[...]) unpacked_fp6_out = self.pallas_call( unpack_kernel, out_shape=jax.ShapeDtypeStruct(shape, jnp.int8), )(two_bits_in, four_bits_in) # Validate that unpacking produces valid 6-bit containers (bits [7:6] are 0) np.testing.assert_array_equal(unpacked_fp6_out.astype(jnp.uint8) >> 6, 0) expected_unpacked = ( ((two_bits_in.astype(jnp.uint8) & 0x3) << 4) | (four_bits_in.astype(jnp.uint8) & 0xF) ).astype(jnp.int8) np.testing.assert_array_equal(unpacked_fp6_out, expected_unpacked) def test_fp6_pack(self): if not jtu.is_device_tpu_at_least(version=5): self.skipTest("Requires TPU v5+") shape = (256, 128) fp6_input = jax.random.randint( jax.random.key(0), shape, 0, 64, dtype=jnp.uint8 ) def pack_kernel(fp6_ref, two_b_out_ref, four_b_out_ref): two_b, four_b = pltpu.pack_fp6(fp6_ref[...]) two_b_out_ref[...] = two_b four_b_out_ref[...] = four_b two_bits_out, four_bits_out = self.pallas_call( pack_kernel, out_shape=( jax.ShapeDtypeStruct(shape, jnp.int2), jax.ShapeDtypeStruct(shape, jnp.int4), ), )(fp6_input) expected_four_b = (fp6_input & 0xF).astype(jnp.int4) expected_two_b = ((fp6_input >> 4) & 0x3).astype(jnp.int2) np.testing.assert_array_equal(four_bits_out, expected_four_b) np.testing.assert_array_equal(two_bits_out, expected_two_b) def test_fp6_layout_roundtrip(self): if not jtu.is_device_tpu_at_least(version=5): self.skipTest("Requires TPU v5+") shape = (256, 128) two_bits_in = jax.random.randint( jax.random.key(0), shape, -2, 2, dtype=jnp.int2 ) four_bits_in = jax.random.randint( jax.random.key(1), shape, -8, 8, dtype=jnp.int4 ) def unpack_kernel(two_b_ref, four_b_ref, o_ref): o_ref[...] = pltpu.unpack_fp6(two_b_ref[...], four_b_ref[...]) unpacked_fp6_out = self.pallas_call( unpack_kernel, out_shape=jax.ShapeDtypeStruct(shape, jnp.int8), )(two_bits_in, four_bits_in) def pack_kernel(fp6_ref, two_b_out_ref, four_b_out_ref): two_b, four_b = pltpu.pack_fp6(fp6_ref[...]) two_b_out_ref[...] = two_b four_b_out_ref[...] = four_b two_bits_out, four_bits_out = self.pallas_call( pack_kernel, out_shape=( jax.ShapeDtypeStruct(shape, jnp.int2), jax.ShapeDtypeStruct(shape, jnp.int4), ), )(unpacked_fp6_out) np.testing.assert_array_equal(two_bits_out, two_bits_in) np.testing.assert_array_equal(four_bits_out, four_bits_in) class ConvTest(ptu.PallasTPUTest): def setUp(self): super().setUp() if not jtu.is_device_tpu_at_least(version=4): self.skipTest("Requires TPUv4+") if not jtu.is_libtpu_at_least("0.1.0"): self.skipTest("Mosaic TPU conv requires libtpu >= 0.1.0") def test_conv_general_dilated_1d(self): x = jax.random.normal(jax.random.key(0), (2, 16, 8), dtype=jnp.float32) kernel = jax.random.normal(jax.random.key(1), (3, 8, 4), dtype=jnp.float32) def conv_kernel(x_ref, k_ref, o_ref): o_ref[...] = lax.conv_general_dilated( x_ref[...], k_ref[...], window_strides=(1,), padding=((1, 1),), dimension_numbers=("NHC", "HIO", "NHC"), ) out = pl.pallas_call( conv_kernel, out_shape=jax.ShapeDtypeStruct((2, 16, 4), jnp.float32), )(x, kernel) expected = lax.conv_general_dilated( x, kernel, window_strides=(1,), padding=((1, 1),), dimension_numbers=("NHC", "HIO", "NHC"), ) self.assertAllClose(out, expected, rtol=1e-5, atol=1e-5) @parameterized.named_parameters( dict( testcase_name="_basic", lhs_shape=(1, 8, 8, 8), rhs_shape=(3, 3, 8, 4), window_strides=(1, 1), padding=((0, 0), (0, 0)), ), dict( testcase_name="_rhs_non_contracting_spatial", lhs_shape=(8, 1, 1, 128), rhs_shape=(4, 8, 128, 128), window_strides=(1, 1), padding=((3, 3), (7, 7)), window_reversal=(True, True), ), dict( testcase_name="_padded_batch_spatial", lhs_shape=(8, 4, 8, 128), rhs_shape=(4, 8, 128, 128), window_strides=(4, 8), padding=((3, 3), (7, 7)), lhs_dilation=(3, 7), window_reversal=(True, True), ), ) def test_pltpu_conv_2d( self, *, lhs_shape, rhs_shape, window_strides, padding, lhs_dilation=None, window_reversal=None, ): x, kernel = self._conv_operands(lhs_shape, rhs_shape, jnp.float32) def conv_kernel(x_ref, k_ref, o_ref): o_ref[...] = pltpu.conv( x_ref[...], k_ref[...], window_strides=window_strides, padding=padding, lhs_dilation=lhs_dilation, window_reversal=window_reversal, dimension_numbers=("NHWC", "HWIO", "NHWC"), ) ref_kernel = jnp.flip(kernel, axis=(0, 1)) if window_reversal else kernel expected = lax.conv_general_dilated( x, ref_kernel, window_strides=window_strides, padding=padding, lhs_dilation=lhs_dilation, dimension_numbers=("NHWC", "HWIO", "NHWC"), ) out = pl.pallas_call( conv_kernel, out_shape=jax.ShapeDtypeStruct(expected.shape, jnp.float32), )(x, kernel) self.assertAllClose(out, expected, rtol=1e-5, atol=1e-5) def test_conv_with_strides(self): x = jax.random.normal(jax.random.key(4), (1, 10, 10, 8), dtype=jnp.float32) kernel = jax.random.normal( jax.random.key(5), (2, 2, 8, 4), dtype=jnp.float32 ) def conv_kernel(x_ref, k_ref, o_ref): o_ref[...] = pltpu.conv( x_ref[...], k_ref[...], window_strides=(2, 2), padding=((0, 0), (0, 0)), dimension_numbers=("NHWC", "HWIO", "NHWC"), ) out = pl.pallas_call( conv_kernel, out_shape=jax.ShapeDtypeStruct((1, 5, 5, 4), jnp.float32), )(x, kernel) expected = lax.conv_general_dilated( x, kernel, window_strides=(2, 2), padding="VALID", dimension_numbers=("NHWC", "HWIO", "NHWC"), ) self.assertAllClose(out, expected, rtol=1e-5, atol=1e-5) def test_conv_with_preferred_element_type(self): x, kernel = self._conv_operands( (1, 8, 8, 8), (3, 3, 8, 4), jnp.bfloat16, seed=6 ) def conv_kernel(x_ref, k_ref, o_ref): o_ref[...] = pltpu.conv( x_ref[...], k_ref[...], window_strides=(1, 1), padding=((1, 1), (1, 1)), dimension_numbers=("NHWC", "HWIO", "NHWC"), preferred_element_type=jnp.float32, ) out = pl.pallas_call( conv_kernel, out_shape=jax.ShapeDtypeStruct((1, 8, 8, 4), jnp.float32), )(x, kernel) expected = lax.conv_general_dilated( x, kernel, window_strides=(1, 1), padding="SAME", dimension_numbers=("NHWC", "HWIO", "NHWC"), preferred_element_type=jnp.float32, ) np.testing.assert_array_equal(out, expected) def _conv_operands(self, lhs_shape, rhs_shape, dtype, seed=1234): """Small integer-valued operands, so bf16 MXU passes stay exact.""" k1, k2 = jax.random.split(jax.random.key(seed)) lhs = jax.random.randint( k1, lhs_shape, minval=-4, maxval=5, dtype=jnp.int32 ).astype(dtype) rhs = jax.random.randint( k2, rhs_shape, minval=-4, maxval=5, dtype=jnp.int32 ).astype(dtype) return lhs, rhs @parameterized.named_parameters( dict( testcase_name="_1x1", lhs_shape=(1, 8, 8, 128), rhs_shape=(1, 1, 128, 256), window_strides=(1, 1), padding=((0, 0), (0, 0)), ), dict( testcase_name="_3x3_same", lhs_shape=(1, 8, 8, 128), rhs_shape=(3, 3, 128, 128), window_strides=(1, 1), padding=((1, 1), (1, 1)), ), dict( testcase_name="_3x3_valid", lhs_shape=(1, 10, 10, 128), rhs_shape=(3, 3, 128, 128), window_strides=(1, 1), padding=((0, 0), (0, 0)), ), dict( testcase_name="_3x3_strided", lhs_shape=(1, 16, 16, 128), rhs_shape=(3, 3, 128, 128), window_strides=(2, 2), padding=((1, 1), (1, 1)), ), dict( testcase_name="_2x2_atrous", lhs_shape=(1, 8, 8, 128), rhs_shape=(2, 2, 128, 128), window_strides=(1, 1), padding=((0, 0), (0, 0)), rhs_dilation=(2, 2), ), dict( testcase_name="_2x2_transposed", lhs_shape=(1, 4, 4, 128), rhs_shape=(2, 2, 128, 128), window_strides=(1, 1), padding=((1, 1), (1, 1)), lhs_dilation=(2, 2), ), dict( testcase_name="_3x5_transposed", lhs_shape=(1, 4, 4, 128), rhs_shape=(3, 3, 128, 128), window_strides=(1, 1), padding=((1, 1), (1, 1)), lhs_dilation=(3, 5), ), dict( testcase_name="_3x3_negative_padding", lhs_shape=(1, 12, 12, 128), rhs_shape=(3, 3, 128, 128), window_strides=(1, 1), padding=((-1, -1), (-1, 2)), ), dict( testcase_name="_batched", lhs_shape=(4, 8, 8, 128), rhs_shape=(3, 3, 128, 128), window_strides=(1, 1), padding=((1, 1), (1, 1)), ), dict( testcase_name="_feature_grouped_2x", lhs_shape=(1, 8, 8, 128), rhs_shape=(3, 3, 64, 128), window_strides=(1, 1), padding=((1, 1), (1, 1)), feature_group_count=2, ), dict( testcase_name="_feature_grouped_4x", lhs_shape=(1, 8, 8, 128), rhs_shape=(3, 3, 32, 128), window_strides=(1, 1), padding=((1, 1), (1, 1)), feature_group_count=4, ), dict( testcase_name="_depthwise", lhs_shape=(1, 8, 8, 16), rhs_shape=(3, 3, 1, 16), window_strides=(1, 1), padding=((1, 1), (1, 1)), feature_group_count=16, ), dict( testcase_name="_batch_grouped_2x", lhs_shape=(4, 8, 8, 128), rhs_shape=(3, 3, 128, 64), window_strides=(1, 1), padding=((1, 1), (1, 1)), batch_group_count=2, ), dict( testcase_name="_mixed_degenerate", lhs_shape=(8, 1, 16, 128), rhs_shape=(1, 3, 128, 128), window_strides=(1, 1), padding=((0, 0), (1, 1)), ), dict( testcase_name="_lhs_non_contracting_spatial", lhs_shape=(2, 1, 8, 128), rhs_shape=(1, 1, 128, 128), window_strides=(1, 1), padding=((0, 0), (0, 0)), ), dict( testcase_name="_unpadded_batch_spatial", lhs_shape=(8, 4, 8, 128), rhs_shape=(4, 8, 128, 128), window_strides=(3, 7), padding=((0, 0), (0, 0)), lhs_dilation=(4, 8), ), ) def test_conv_general_dilated( self, *, lhs_shape, rhs_shape, window_strides, padding, lhs_dilation=None, rhs_dilation=None, feature_group_count=1, batch_group_count=1, ): dimension_numbers = ("NHWC", "HWIO", "NHWC") conv = functools.partial( jax.lax.conv_general_dilated, window_strides=window_strides, padding=padding, lhs_dilation=lhs_dilation, rhs_dilation=rhs_dilation, dimension_numbers=dimension_numbers, feature_group_count=feature_group_count, batch_group_count=batch_group_count, ) def kernel(x_ref, w_ref, o_ref): o_ref[...] = conv(x_ref[...], w_ref[...]) lhs, rhs = self._conv_operands(lhs_shape, rhs_shape, jnp.float32) expected = conv(lhs, rhs) out = jax.ShapeDtypeStruct(expected.shape, expected.dtype) result = self.pallas_call(kernel, out_shape=out)(lhs, rhs) np.testing.assert_array_equal(result, expected) @parameterized.named_parameters( dict( testcase_name="_nchw_hwio_nhwc", dimension_numbers=("NCHW", "HWIO", "NHWC"), lhs_shape=(1, 128, 8, 8), rhs_shape=(3, 3, 128, 128), ), dict( testcase_name="_nhwc_ihwo_nhwc", dimension_numbers=("NHWC", "IHWO", "NHWC"), lhs_shape=(1, 8, 8, 128), rhs_shape=(128, 3, 3, 128), ), ) def test_conv_dimension_numbers( self, *, dimension_numbers, lhs_shape, rhs_shape ): conv = functools.partial( jax.lax.conv_general_dilated, window_strides=(1, 1), padding="SAME", dimension_numbers=dimension_numbers, ) def kernel(x_ref, w_ref, o_ref): o_ref[...] = conv(x_ref[...], w_ref[...]) lhs, rhs = self._conv_operands(lhs_shape, rhs_shape, jnp.float32) expected = conv(lhs, rhs) out = jax.ShapeDtypeStruct(expected.shape, expected.dtype) result = self.pallas_call(kernel, out_shape=out)(lhs, rhs) np.testing.assert_array_equal(result, expected) def test_conv_0d(self): conv = functools.partial( jax.lax.conv_general_dilated, window_strides=(), padding=(), dimension_numbers=("NC", "IO", "NC"), ) def kernel(x_ref, w_ref, o_ref): o_ref[...] = conv(x_ref[...], w_ref[...]) lhs, rhs = self._conv_operands((8, 128), (128, 128), jnp.float32) expected = conv(lhs, rhs) out = jax.ShapeDtypeStruct(expected.shape, expected.dtype) result = self.pallas_call(kernel, out_shape=out)(lhs, rhs) np.testing.assert_array_equal(result, expected) def test_conv_1d(self): conv = functools.partial( jax.lax.conv_general_dilated, window_strides=(2,), padding=((1, 1),), rhs_dilation=(2,), dimension_numbers=("NWC", "WIO", "NWC"), ) def kernel(x_ref, w_ref, o_ref): o_ref[...] = conv(x_ref[...], w_ref[...]) lhs, rhs = self._conv_operands((8, 16, 128), (3, 128, 128), jnp.float32) expected = conv(lhs, rhs) out = jax.ShapeDtypeStruct(expected.shape, expected.dtype) result = self.pallas_call(kernel, out_shape=out)(lhs, rhs) np.testing.assert_array_equal(result, expected) def test_conv_3d(self): conv = functools.partial( jax.lax.conv_general_dilated, window_strides=(1, 1, 1), padding="SAME", dimension_numbers=("NDHWC", "DHWIO", "NDHWC"), ) def kernel(x_ref, w_ref, o_ref): o_ref[...] = conv(x_ref[...], w_ref[...]) lhs, rhs = self._conv_operands( (1, 2, 4, 8, 128), (2, 3, 3, 128, 128), jnp.float32 ) expected = conv(lhs, rhs) out = jax.ShapeDtypeStruct(expected.shape, expected.dtype) result = self.pallas_call(kernel, out_shape=out)(lhs, rhs) np.testing.assert_array_equal(result, expected) @parameterized.parameters( (jnp.float32, None), (jnp.bfloat16, jnp.float32), (jnp.float32, jnp.float32), ) def test_conv_dtypes(self, dtype, preferred_element_type): conv = functools.partial( jax.lax.conv_general_dilated, window_strides=(1, 1), padding="SAME", dimension_numbers=("NHWC", "HWIO", "NHWC"), preferred_element_type=preferred_element_type, ) def kernel(x_ref, w_ref, o_ref): o_ref[...] = conv(x_ref[...], w_ref[...]) lhs, rhs = self._conv_operands((1, 8, 8, 128), (3, 3, 128, 128), dtype) expected = conv(lhs, rhs) out = jax.ShapeDtypeStruct(expected.shape, expected.dtype) result = self.pallas_call(kernel, out_shape=out)(lhs, rhs) self.assertEqual(result.dtype, expected.dtype) np.testing.assert_array_equal(result, expected) def _pallas_dot_supports(self, dtype, preferred_element_type): """Whether a Pallas dot_general accepts these dtypes on this chip.""" def kernel(x_ref, w_ref, o_ref): o_ref[...] = jax.lax.dot_general( x_ref[...], w_ref[...], (((1,), (0,)), ((), ())), preferred_element_type=preferred_element_type, ) lhs, rhs = self._conv_operands((128, 128), (128, 128), dtype) out = jax.ShapeDtypeStruct( (128, 128), jnp.dtype(preferred_element_type or dtype) ) try: self.pallas_call(kernel, out_shape=out)(lhs, rhs) except Exception: return False return True # A convolution lowers to a chain of tpu.matmul, so it must accept exactly # the dtype combinations a Pallas dot_general accepts on the same chip. The # test probes dot_general first and only skips when it is unsupported too, so # any combination where conv is more restrictive than matmul fails here. @parameterized.named_parameters( dict( testcase_name="_int8", dtype=jnp.int8, preferred_element_type=jnp.int32, ), dict( testcase_name="_int8_f32_acc", dtype=jnp.int8, preferred_element_type=jnp.float32, ), dict( testcase_name="_int4", dtype=jnp.int4, preferred_element_type=jnp.int32, ), dict( testcase_name="_f8e5m2", dtype=jnp.float8_e5m2, preferred_element_type=jnp.float32, ), dict( testcase_name="_f8e4m3fn", dtype=jnp.float8_e4m3fn, preferred_element_type=jnp.float32, ), ) def test_conv_narrow_dtypes(self, *, dtype, preferred_element_type): if not jtu.is_device_tpu_at_least(version=5): self.skipTest("TPUv5+ needed for narrow dtype matmuls") if not self._pallas_dot_supports(dtype, preferred_element_type): self.skipTest( f"dot_general does not support {dtype.__name__} with a" f" {preferred_element_type.__name__} accumulator on this chip either" ) conv = functools.partial( jax.lax.conv_general_dilated, window_strides=(1, 1), padding="SAME", dimension_numbers=("NHWC", "HWIO", "NHWC"), preferred_element_type=preferred_element_type, ) def kernel(x_ref, w_ref, o_ref): o_ref[...] = conv(x_ref[...], w_ref[...]) lhs, rhs = self._conv_operands((1, 8, 8, 128), (3, 3, 128, 128), dtype) expected = conv(lhs, rhs) out = jax.ShapeDtypeStruct(expected.shape, expected.dtype) result = self.pallas_call(kernel, out_shape=out)(lhs, rhs) self.assertEqual(result.dtype, expected.dtype) np.testing.assert_array_equal(result, expected) def test_conv_unsupported_dtype_reports_an_error(self): # An unsupported operand dtype must surface as a lowering error rather than # crashing the compiler inside canonicalize-mosaic. if not jtu.is_device_tpu_at_least(version=5): self.skipTest("TPUv5+ needed for narrow dtype matmuls") if self._pallas_dot_supports(jnp.int8, jnp.float32): self.skipTest("This chip supports an int8 matmul with a float32 acc") conv = functools.partial( jax.lax.conv_general_dilated, window_strides=(1, 1), padding="SAME", dimension_numbers=("NHWC", "HWIO", "NHWC"), preferred_element_type=jnp.float32, ) def kernel(x_ref, w_ref, o_ref): o_ref[...] = conv(x_ref[...], w_ref[...]) lhs, rhs = self._conv_operands((1, 8, 8, 128), (3, 3, 128, 128), jnp.int8) out = jax.ShapeDtypeStruct((1, 8, 8, 128), jnp.float32) with self.assertRaises(Exception): self.pallas_call(kernel, out_shape=out)(lhs, rhs) @parameterized.parameters( jax.lax.Precision.DEFAULT, jax.lax.Precision.HIGHEST, ) def test_conv_precision(self, precision): conv = functools.partial( jax.lax.conv_general_dilated, window_strides=(1, 1), padding="SAME", dimension_numbers=("NHWC", "HWIO", "NHWC"), precision=precision, ) def kernel(x_ref, w_ref, o_ref): o_ref[...] = conv(x_ref[...], w_ref[...]) lhs, rhs = self._conv_operands( (1, 8, 8, 128), (3, 3, 128, 128), jnp.float32 ) expected = conv(lhs, rhs) out = jax.ShapeDtypeStruct(expected.shape, expected.dtype) result = self.pallas_call(kernel, out_shape=out)(lhs, rhs) np.testing.assert_array_equal(result, expected) def test_pltpu_conv(self): def kernel(x_ref, w_ref, o_ref): o_ref[...] = pltpu.conv( x_ref[...], w_ref[...], window_strides=(1, 1), padding="SAME", dimension_numbers=("NHWC", "HWIO", "NHWC"), ) lhs, rhs = self._conv_operands( (1, 8, 8, 128), (3, 3, 128, 128), jnp.float32 ) expected = jax.lax.conv_general_dilated( lhs, rhs, window_strides=(1, 1), padding="SAME", dimension_numbers=("NHWC", "HWIO", "NHWC"), ) out = jax.ShapeDtypeStruct(expected.shape, expected.dtype) result = self.pallas_call(kernel, out_shape=out)(lhs, rhs) np.testing.assert_array_equal(result, expected) @parameterized.parameters(jnp.uint32, jnp.int32) def test_mulhi(self, dtype): if not jtu.is_libtpu_at_least("0.0.51"): self.skipTest("Requires libtpu >= 0.0.51") shape = (8, 128) rng = np.random.default_rng(42) if dtype == jnp.uint32: x_np = rng.integers(0, 2**32 - 1, size=shape, dtype=np.uint32, endpoint=True) y_np = rng.integers(0, 2**32 - 1, size=shape, dtype=np.uint32, endpoint=True) x_np.ravel()[:4] = [0, 1, 2**32 - 1, 2**31] y_np.ravel()[:4] = [2**32 - 1, 2**32 - 1, 2**32 - 1, 2] else: x_np = rng.integers(-(2**31), 2**31 - 1, size=shape, dtype=np.int32, endpoint=True) y_np = rng.integers(-(2**31), 2**31 - 1, size=shape, dtype=np.int32, endpoint=True) x_np.ravel()[:4] = [0, 1, -1, -(2**31)] y_np.ravel()[:4] = [-(2**31), 2**31 - 1, -1, -(2**31)] x = jnp.asarray(x_np) y = jnp.asarray(y_np) def hi_kernel(x_ref, y_ref, hi_ref): hi_ref[...] = jax.lax.mulhi(x_ref[...], y_ref[...]) out_spec = jax.ShapeDtypeStruct(shape, dtype) hi_got = self.pallas_call(hi_kernel, out_shape=out_spec)(x, y) np.testing.assert_array_equal(hi_got, jax.lax.mulhi(x, y)) def hilo_kernel(x_ref, y_ref, lo_ref, hi_ref): a, b = x_ref[...], y_ref[...] lo_ref[...] = a * b hi_ref[...] = jax.lax.mulhi(a, b) lo_got, hi_got = self.pallas_call( hilo_kernel, out_shape=(out_spec, out_spec) )(x, y) np.testing.assert_array_equal(lo_got, x * y) np.testing.assert_array_equal(hi_got, jax.lax.mulhi(x, y)) if __name__ == "__main__": absltest.main(testLoader=jtu.JaxTestLoader())