[pallas:mosaic_gpu] Use jax.tree_util.register_dataclass for transforms

PiperOrigin-RevId: 702733084

Author: superbobry

jax/_src/pallas/mosaic_gpu/core.py

@@ -104,7 +104,6 @@ def __str__(self) -> str:
     return self.value
 
   def __call__(
-
       self,
       shape: tuple[int, ...],
       dtype: jnp.dtype,
@@ -161,7 +160,6 @@ class TilingTransform(MemoryRefTransform):
   shape of (M // X, N // Y, X, Y). Ex. A (256, 256) block that is tiled with a
   tiling of (64, 32) will be tiled as (4, 8, 64, 32).
   """
-
   tiling: tuple[int, ...]
 
   def undo(self, ref: pallas_core.TransformedRef) -> pallas_core.TransformedRef:
@@ -176,10 +174,10 @@ def to_gpu_transform(self) -> mgpu.MemRefTransform:
     return mgpu.TileTransform(self.tiling)
 
 
-@tree_util.register_pytree_node_class
+@tree_util.register_dataclass
 @dataclasses.dataclass(frozen=True)
 class UntileRef(state_types.Transform):
-  tiling: tuple[int, ...]
+  tiling: tuple[int, ...] = dataclasses.field(metadata=dict(static=True))
 
   def transform_shape(self, shape):
     if shape is None:
@@ -214,14 +212,6 @@ def untransform_index(
   def undo_to_gpu_transform(self) -> mgpu.MemRefTransform:
     return mgpu.TileTransform(self.tiling)
 
-  def tree_flatten(self):
-    return (), (self.tiling,)
-
-  @classmethod
-  def tree_unflatten(cls, metadata, arrays):
-    assert not arrays
-    return cls(*metadata)
-
 
 def _perm_inverse(permutation: tuple[int, ...]) -> tuple[int, ...]:
   inverse = [-1] * len(permutation)
@@ -257,7 +247,7 @@ def to_gpu_transform(self) -> mgpu.MemRefTransform:
     return mgpu.TransposeTransform(self.permutation)
 
 
-@tree_util.register_pytree_node_class
+@tree_util.register_dataclass
 @dataclasses.dataclass(frozen=True)
 class TransposeRef(state_types.Transform):
   permutation: tuple[int, ...]
@@ -287,14 +277,6 @@ def untransform_index(
   def undo_to_gpu_transform(self) -> mgpu.MemRefTransform:
     return mgpu.TransposeTransform(_perm_inverse(self.permutation))
 
-  def tree_flatten(self):
-    return (), (self.permutation,)
-
-  @classmethod
-  def tree_unflatten(cls, metadata, arrays):
-    assert not arrays
-    return cls(*metadata)
-
 
 def transpose_ref(
     ref: pallas_core.TransformedRef | Any,
@@ -345,10 +327,10 @@ def __call__(self, aval: jax_core.ShapedArray) -> jax_core.ShapedArray:
     return aval
 
 
-@tree_util.register_pytree_node_class
+@tree_util.register_dataclass
 @dataclasses.dataclass(frozen=True)
 class UnswizzleRef(state_types.Transform):
-  swizzle: int
+  swizzle: int = dataclasses.field(metadata=dict(static=True))
 
   def untransform_index(
       self, idxs: tuple[Index, ...]
@@ -369,14 +351,6 @@ def untransform_index(
       raise ValueError("Swizzled dims cannot be sliced")
     return idxs, self
 
-  def tree_flatten(self):
-    return (), (self.swizzle,)
-
-  @classmethod
-  def tree_unflatten(cls, metadata, arrays):
-    assert not arrays
-    return cls(*metadata)
-
 
 @dataclasses.dataclass
 class GPUBlockSpec(pallas_core.BlockSpec):