fixed import errors in test

Author: silkyarora

keras_rs/src/layers/embedding/jax/checkpoint_utils.py

@@ -1,10 +1,10 @@
 """A Wrapper over orbax CheckpointManager for Keras3 Jax TPU Embeddings."""
 
 from typing import Any
-from etils import epath
-import jax
+
 import keras
 import orbax.checkpoint as ocp
+from etils import epath
 
 
 class JaxKeras3CheckpointManager(ocp.CheckpointManager):
@@ -47,20 +47,18 @@ def save_state(self, epoch: int) -> None:
         state, metrics = self._get_state()
         self.save(
             epoch * self._steps_per_epoch,
-            args=ocp.args.StandardSave(state),
-            metrics=ocp.args.StandardSave(metrics),
+            args=ocp.args.StandardSave(item=state),
         )
 
     def restore_state(self, step: int | None = None) -> None:
         """Restores the model from the checkpoint directory.
 
         Args:
-          step: The step number to restore the state from. Default=None
+          step: The step .number to restore the state from. Default=None
             restores the latest step.
         """
         if step is None:
             step = self.latest_step()
-
         # Restore the model state only, not metrics.
         state, _ = self._get_state()
         restored_state = self.restore(

keras_rs/src/layers/embedding/jax/distributed_embedding.py

@@ -384,7 +384,7 @@ def sparsecore_build(
     ) -> None:
         del input_shapes  # Unused.
 
-        if self._sparsecore_built:
+        if hasattr(self, "_sparsecore_built") and self._sparsecore_built:
             return
 
         feature_specs = config_conversion.keras_to_jte_feature_configs(

keras_rs/src/layers/embedding/jax/distributed_embedding_test.py

@@ -19,8 +19,8 @@
 from jax_tpu_embedding.sparsecore.utils import utils as jte_utils
 
 from keras_rs.src.layers.embedding import test_utils as keras_test_utils
-from keras_rs.src.layers.embedding.jax import config_conversion
 from keras_rs.src.layers.embedding.jax import checkpoint_utils
+from keras_rs.src.layers.embedding.jax import config_conversion
 from keras_rs.src.layers.embedding.jax import (
     distributed_embedding as jax_distributed_embedding,
 )
@@ -564,7 +564,10 @@ def loss_fn(y_true, y_pred):
             table_stacking=target_stacking,
             name=embedding_layer_name,
         )
-        layer_2._sparsecore_build()
+        input_shapes = jax.tree.map(
+            lambda f: f.input_shape, feature_configs_with_zero_init
+        )
+        layer_2.build(input_shapes)
         model_for_restore = keras.Sequential([layer_2])
         manager_for_restore = checkpoint_utils.JaxKeras3CheckpointManager(
             model_for_restore,
@@ -574,6 +577,7 @@ def loss_fn(y_true, y_pred):
         )
         model_for_restore.compile(jit_compile=True, loss=loss_fn)
         model_for_restore.build()
+        model_for_restore.optimizer.build(model_for_restore.trainable_variables)
         manager_for_restore.restore_state()
         jax.tree.map(
             np.testing.assert_array_equal,