keras-rs-nightly 0.2.2.dev202506170335__tar.gz → 0.2.2.dev202506190335__tar.gz

{keras_rs_nightly-0.2.2.dev202506170335 → keras_rs_nightly-0.2.2.dev202506190335}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: keras-rs-nightly
-Version: 0.2.2.dev202506170335
+Version: 0.2.2.dev202506190335
 Summary: Multi-backend recommender systems with Keras 3.
 Author-email: Keras team <keras-users@googlegroups.com>
 License: Apache License 2.0

{keras_rs_nightly-0.2.2.dev202506170335 → keras_rs_nightly-0.2.2.dev202506190335}/keras_rs/src/layers/embedding/base_distributed_embedding.py

@@ -174,12 +174,9 @@ class DistributedEmbedding(keras.layers.Layer):
     supported on all backends and accelerators:
 
     - `keras.optimizers.Adagrad`
-    - `keras.optimizers.SGD`
-
-    The following are additionally available when using the TensorFlow backend:
-
     - `keras.optimizers.Adam`
     - `keras.optimizers.Ftrl`
+    - `keras.optimizers.SGD`
 
     Also, not all parameters of the optimizers are supported (e.g. the
     `nesterov` option of `SGD`). An error is raised when an unsupported

keras_rs_nightly-0.2.2.dev202506190335/keras_rs/src/layers/embedding/jax/checkpoint_utils.py

@@ -0,0 +1,104 @@
+"""A Wrapper over orbax CheckpointManager for Keras3 Jax TPU Embeddings."""
+
+from typing import Any
+
+import keras
+import orbax.checkpoint as ocp
+from etils import epath
+
+
+class JaxKeras3CheckpointManager(ocp.CheckpointManager):
+    """A wrapper over orbax CheckpointManager for Keras3 Jax TPU Embeddings."""
+
+    def __init__(
+        self,
+        model: keras.Model,
+        checkpoint_dir: epath.PathLike,
+        max_to_keep: int,
+        steps_per_epoch: int = 1,
+        **kwargs: Any,
+    ):
+        options = ocp.CheckpointManagerOptions(
+            max_to_keep=max_to_keep, enable_async_checkpointing=False, **kwargs
+        )
+        self._model = model
+        self._steps_per_epoch = steps_per_epoch
+        self._checkpoint_dir = checkpoint_dir
+        super().__init__(checkpoint_dir, options=options)
+
+    def _get_state(self) -> tuple[dict[str, Any], Any | None]:
+        """Gets the model state and metrics"""
+        model_state = self._model.get_state_tree()
+        state = {}
+        metrics = None
+        for k, v in model_state.items():
+            if k == "metrics_variables":
+                metrics = v
+            else:
+                state[k] = v
+        return state, metrics
+
+    def save_state(self, epoch: int) -> None:
+        """Saves the model to the checkpoint directory.
+
+        Args:
+          epoch: The epoch number at which the state is saved.
+        """
+        state, metrics_value = self._get_state()
+        self.save(
+            epoch * self._steps_per_epoch,
+            args=ocp.args.StandardSave(item=state),
+            metrics=metrics_value,
+        )
+
+    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
+            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(
+            step, args=ocp.args.StandardRestore(item=state)
+        )
+        self._model.set_state_tree(restored_state)
+
+
+class JaxKeras3CheckpointCallback(keras.callbacks.Callback):
+    """A callback for checkpointing and restoring state using Orbax."""
+
+    def __init__(
+        self,
+        model: keras.Model,
+        checkpoint_dir: epath.PathLike,
+        max_to_keep: int,
+        steps_per_epoch: int = 1,
+        **kwargs: Any,
+    ):
+        if keras.backend.backend() != "jax":
+            raise ValueError(
+                "`JaxKeras3CheckpointCallback` is only supported on a "
+                "`jax` backend."
+            )
+        self._checkpoint_manager = JaxKeras3CheckpointManager(
+            model, checkpoint_dir, max_to_keep, steps_per_epoch, **kwargs
+        )
+
+    def on_train_begin(self, logs: dict[str, Any] | None = None) -> None:
+        if not self.model.built or not self.model.optimizer.built:
+            raise ValueError(
+                "To use `JaxKeras3CheckpointCallback`, your model and "
+                "optimizer must be built before you call `fit()`."
+            )
+        latest_epoch = self._checkpoint_manager.latest_step()
+        if latest_epoch is not None:
+            self._checkpoint_manager.restore_state(step=latest_epoch)
+
+    def on_epoch_end(
+        self, epoch: int, logs: dict[str, Any] | None = None
+    ) -> None:
+        self._checkpoint_manager.save_state(epoch)

{keras_rs_nightly-0.2.2.dev202506170335 → keras_rs_nightly-0.2.2.dev202506190335}/keras_rs/src/layers/embedding/jax/config_conversion.py

@@ -229,18 +229,63 @@ def keras_to_jte_optimizer(
     # pylint: disable-next=protected-access
     learning_rate = keras_to_jte_learning_rate(optimizer._learning_rate)
 
-    # SGD or Adagrad
+    # Unsupported keras optimizer general options.
+    if optimizer.clipnorm is not None:
+        raise ValueError("Unsupported optimizer option `clipnorm`.")
+    if optimizer.global_clipnorm is not None:
+        raise ValueError("Unsupported optimizer option `global_clipnorm`.")
+    if optimizer.use_ema:
+        raise ValueError("Unsupported optimizer option `use_ema`.")
+    if optimizer.loss_scale_factor is not None:
+        raise ValueError("Unsupported optimizer option `loss_scale_factor`.")
+
+    # Supported optimizers.
     if isinstance(optimizer, keras.optimizers.SGD):
+        if getattr(optimizer, "nesterov", False):
+            raise ValueError("Unsupported optimizer option `nesterov`.")
+        if getattr(optimizer, "momentum", 0.0) != 0.0:
+            raise ValueError("Unsupported optimizer option `momentum`.")
         return embedding_spec.SGDOptimizerSpec(learning_rate=learning_rate)
     elif isinstance(optimizer, keras.optimizers.Adagrad):
+        if getattr(optimizer, "epsilon", 1e-7) != 1e-7:
+            raise ValueError("Unsupported optimizer option `epsilon`.")
         return embedding_spec.AdagradOptimizerSpec(
             learning_rate=learning_rate,
             initial_accumulator_value=optimizer.initial_accumulator_value,
         )
+    elif isinstance(optimizer, keras.optimizers.Adam):
+        if getattr(optimizer, "amsgrad", False):
+            raise ValueError("Unsupported optimizer option `amsgrad`.")
 
-    # Default to SGD for now, since other optimizers are still being created,
-    # and we don't want to fail.
-    return embedding_spec.SGDOptimizerSpec(learning_rate=learning_rate)
+        return embedding_spec.AdamOptimizerSpec(
+            learning_rate=learning_rate,
+            beta_1=optimizer.beta_1,
+            beta_2=optimizer.beta_2,
+            epsilon=optimizer.epsilon,
+        )
+    elif isinstance(optimizer, keras.optimizers.Ftrl):
+        if (
+            getattr(optimizer, "l2_shrinkage_regularization_strength", 0.0)
+            != 0.0
+        ):
+            raise ValueError(
+                "Unsupported optimizer option "
+                "`l2_shrinkage_regularization_strength`."
+            )
+
+        return embedding_spec.FTRLOptimizerSpec(
+            learning_rate=learning_rate,
+            learning_rate_power=optimizer.learning_rate_power,
+            l1_regularization_strength=optimizer.l1_regularization_strength,
+            l2_regularization_strength=optimizer.l2_regularization_strength,
+            beta=optimizer.beta,
+            initial_accumulator_value=optimizer.initial_accumulator_value,
+        )
+
+    raise ValueError(
+        f"Unsupported optimizer type {type(optimizer)}. Optimizer must be "
+        f"one of [Adagrad, Adam, Ftrl, SGD]."
+    )
 
 
 def jte_to_keras_optimizer(
@@ -262,8 +307,33 @@ def jte_to_keras_optimizer(
             learning_rate=learning_rate,
             initial_accumulator_value=optimizer.initial_accumulator_value,
         )
+    elif isinstance(optimizer, embedding_spec.AdamOptimizerSpec):
+        return keras.optimizers.Adam(
+            learning_rate=learning_rate,
+            beta_1=optimizer.beta_1,
+            beta_2=optimizer.beta_2,
+            epsilon=optimizer.epsilon,
+        )
+    elif isinstance(optimizer, embedding_spec.FTRLOptimizerSpec):
+        if getattr(optimizer, "initial_linear_value", 0.0) != 0.0:
+            raise ValueError(
+                "Unsupported optimizer option `initial_linear_value`."
+            )
+        if getattr(optimizer, "multiply_linear_by_learning_rate", False):
+            raise ValueError(
+                "Unsupported optimizer option "
+                "`multiply_linear_by_learning_rate`."
+            )
+        return keras.optimizers.Ftrl(
+            learning_rate=learning_rate,
+            learning_rate_power=optimizer.learning_rate_power,
+            initial_accumulator_value=optimizer.initial_accumulator_value,
+            l1_regularization_strength=optimizer.l1_regularization_strength,
+            l2_regularization_strength=optimizer.l2_regularization_strength,
+            beta=optimizer.beta,
+        )
 
-    raise ValueError(f"Unknown optimizer spec {optimizer}")
+    raise ValueError(f"Unknown optimizer spec {type(optimizer)}.")
 
 
 def _keras_to_jte_table_config(

{keras_rs_nightly-0.2.2.dev202506170335 → keras_rs_nightly-0.2.2.dev202506190335}/keras_rs/src/version.py

@@ -1,7 +1,7 @@
 from keras_rs.src.api_export import keras_rs_export
 
 # Unique source of truth for the version number.
-__version__ = "0.2.2.dev202506170335"
+__version__ = "0.2.2.dev202506190335"
 
 
 @keras_rs_export("keras_rs.version")

{keras_rs_nightly-0.2.2.dev202506170335 → keras_rs_nightly-0.2.2.dev202506190335}/keras_rs_nightly.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: keras-rs-nightly
-Version: 0.2.2.dev202506170335
+Version: 0.2.2.dev202506190335
 Summary: Multi-backend recommender systems with Keras 3.
 Author-email: Keras team <keras-users@googlegroups.com>
 License: Apache License 2.0

{keras_rs_nightly-0.2.2.dev202506170335 → keras_rs_nightly-0.2.2.dev202506190335}/keras_rs_nightly.egg-info/SOURCES.txt

@@ -15,6 +15,7 @@ keras_rs/src/layers/embedding/distributed_embedding.py
 keras_rs/src/layers/embedding/distributed_embedding_config.py
 keras_rs/src/layers/embedding/embed_reduce.py
 keras_rs/src/layers/embedding/jax/__init__.py
+keras_rs/src/layers/embedding/jax/checkpoint_utils.py
 keras_rs/src/layers/embedding/jax/config_conversion.py
 keras_rs/src/layers/embedding/jax/distributed_embedding.py
 keras_rs/src/layers/embedding/jax/embedding_lookup.py