[Pallas] Fix 64-bit multiplication emulation in Philox kernel

PiperOrigin-RevId: 819807662

Author: justinjfu

jax/experimental/pallas/ops/tpu/random/philox.py

@@ -46,7 +46,9 @@ def mul32_hi_lo(x: jax.Array, y: jax.Array) -> tuple[jax.Array, jax.Array]:
   cross_xy = xhi * ylo
   cross_yx = xlo * yhi
   carry = (cross_xy & 0xffff) + (cross_yx & 0xffff) + (xy_lo >> 16)
-  return xy_hi + (cross_xy >> 16) + (cross_yx >> 16) + (carry >> 16), xy_lo
+  result_hi = xy_hi + (cross_xy >> 16) + (cross_yx >> 16) + (carry >> 16)
+  result_lo = (carry << 16) + (xy_lo & 0xffff)
+  return result_hi, result_lo
 
 
 def philox_4x32(hi0, lo0, hi1, lo1, k_hi, k_lo, rounds = 10):

tests/pallas/tpu_pallas_random_test.py

@@ -376,6 +376,19 @@ def setUp(self):
       self.skipTest("Need TPU devices")
     super().setUp()
 
+  @parameterized.product(
+      x=[0x1, 0x10000, 0xabcdef],
+      y=[0x1, 0x10000, 0xabcdef],
+  )
+  def test_mul_hi_lo(self, x, y):
+    x = jnp.uint32(x)
+    y = jnp.uint32(y)
+    hi, lo = philox.mul32_hi_lo(x, y)
+    with jax.enable_x64():
+      result = (hi.astype(jnp.uint64) << 32) + lo.astype(jnp.uint64)
+      ref = x.astype(jnp.uint64) * y.astype(jnp.uint64)
+      self.assertEqual(result, ref)
+
   @parameterized.parameters(
       ((512, 512),),
       ((137, 275),),  # Non block-aligned shape