fix-forward for pallas tpu memory spaces test

PiperOrigin-RevId: 770929681

Author: sharadmv

jax/_src/pallas/mosaic/pallas_call_registration.py

@@ -83,6 +83,8 @@ def _get_memory_space_from_aval(
       return None
     case tpu_core.MemorySpace.ANY:
       return None
+    case tpu_core.MemorySpace.HBM:
+      return tpu_custom_call.MemorySpace.HBM
     case tpu_core.MemorySpace.VMEM:
       return tpu_custom_call.MemorySpace.VMEM
     case tpu_core.MemorySpace.SMEM:

jax/_src/pallas/mosaic/primitives.py

@@ -802,6 +802,8 @@ def with_memory_space_constraint(
   Returns:
     The array `x` with the memory space constraint.
   """
+  if memory_space in {tpu_core.MemorySpace.ANY, pl_core.MemorySpace.ANY}:
+    return x
   if memory_space not in {tpu_core.MemorySpace.HBM, tpu_core.MemorySpace.VMEM}:
     raise NotImplementedError(
         "with_memory_space_constraint only supports HBM and VMEM."

tests/pallas/tpu_pallas_memory_space_test.py

@@ -55,7 +55,10 @@ def g(x):
     @jax.jit
     def f(x):
       x = pltpu.with_memory_space_constraint(x, memory_space=memory_space)
-      self.assertEqual(pltpu.get_memory_space(x), memory_space)
+      if color is None:
+        self.assertIsNone(pltpu.get_memory_space(x))
+      else:
+        self.assertEqual(pltpu.get_memory_space(x), memory_space)
       x = g(x)
       return x