Fix for removal of JAX DeviceLocalLayout. (#127)

hertschuh · web-flow · commit a94d9fc473d7 · 2025-06-23T08:57:24.000-07:00
Type verification fails after removal.
diff --git a/keras_rs/src/layers/embedding/jax/distributed_embedding.py b/keras_rs/src/layers/embedding/jax/distributed_embedding.py
@@ -219,7 +219,7 @@ def _create_sparsecore_distribution(
         LayoutClass = (
             jax_layout.Layout
             if jax.__version_info__ >= (0, 6, 3)
-            else jax_layout.DeviceLocalLayout
+            else jax_layout.DeviceLocalLayout  # type: ignore
         )
         # pylint: disable-next=protected-access
         sparsecore_layout._backend_layout = jax_layout.Format(
diff --git a/keras_rs/src/layers/embedding/jax/distributed_embedding_test.py b/keras_rs/src/layers/embedding/jax/distributed_embedding_test.py
@@ -44,7 +44,7 @@ def _create_sparsecore_layout(
     LayoutClass = (
         jax_layout.Layout
         if jax.__version_info__ >= (0, 6, 3)
-        else jax_layout.DeviceLocalLayout
+        else jax_layout.DeviceLocalLayout  # type: ignore
     )
     # pylint: disable-next=protected-access
     sparsecore_layout._backend_layout = jax_layout.Format(

Original file line number	Diff line number	Diff line change
`@@ -219,7 +219,7 @@ def _create_sparsecore_distribution(`
`219`	`219`	`LayoutClass = (`
`220`	`220`	`jax_layout.Layout`
`221`	`221`	`if jax.__version_info__ >= (0, 6, 3)`
`222`		`- else jax_layout.DeviceLocalLayout`
	`222`	`+ else jax_layout.DeviceLocalLayout # type: ignore`
`223`	`223`	`)`
`224`	`224`	`# pylint: disable-next=protected-access`
`225`	`225`	`sparsecore_layout._backend_layout = jax_layout.Format(`
Original file line number	Diff line number	Diff line change
`@@ -44,7 +44,7 @@ def _create_sparsecore_layout(`
`44`	`44`	`LayoutClass = (`
`45`	`45`	`jax_layout.Layout`
`46`	`46`	`if jax.__version_info__ >= (0, 6, 3)`
`47`		`- else jax_layout.DeviceLocalLayout`
	`47`	`+ else jax_layout.DeviceLocalLayout # type: ignore`
`48`	`48`	`)`
`49`	`49`	`# pylint: disable-next=protected-access`
`50`	`50`	`sparsecore_layout._backend_layout = jax_layout.Format(`