Use jax_layout.Format in place of jax_layout.Layout

yashk2810 · recml authors · commit 8e4c5d345d22 · 2025-06-16T11:50:02.000-07:00
JAX is undergoing a rename of the contents of `jax.experimental.layouts` in preparation for its graduation from experimental: 1. "Formats" are replacing "layouts", and specifically `Layout` -> `Format` 2. "Layouts" are replacing "device local layouts", and specifically `DeviceLocalLayout` -> `Layout` This is an incremental update carrying out #1. PiperOrigin-RevId: 772114015
diff --git a/recml/layers/linen/sparsecore.py b/recml/layers/linen/sparsecore.py
@@ -381,7 +381,7 @@ class SparsecoreLayout(nn.Partitioned[A]):
 
   def get_sharding(self, _):
     assert self.mesh is not None
-    return layout.Layout(
+    return layout.Format(
         layout.DeviceLocalLayout(major_to_minor=(0, 1), _tiling=((8,),)),
         jax.sharding.NamedSharding(self.mesh, self.get_partition_spec()),
     )

Original file line number	Diff line number	Diff line change
`@@ -381,7 +381,7 @@ class SparsecoreLayout(nn.Partitioned[A]):`
`381`	`381`
`382`	`382`	`def get_sharding(self, _):`
`383`	`383`	`assert self.mesh is not None`
`384`		`- return layout.Layout(`
	`384`	`+ return layout.Format(`
`385`	`385`	`layout.DeviceLocalLayout(major_to_minor=(0, 1), _tiling=((8,),)),`
`386`	`386`	`jax.sharding.NamedSharding(self.mesh, self.get_partition_spec()),`
`387`	`387`	`)`