keras-rs-nightly 0.2.2.dev202506160338__py3-none-any.whl → 0.2.2.dev202506180335__py3-none-any.whl

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/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/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/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.dev202506160338"
+__version__ = "0.2.2.dev202506180335"
 
 
 @keras_rs_export("keras_rs.version")

{keras_rs_nightly-0.2.2.dev202506160338.dist-info → keras_rs_nightly-0.2.2.dev202506180335.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: keras-rs-nightly
-Version: 0.2.2.dev202506160338
+Version: 0.2.2.dev202506180335
 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.dev202506160338.dist-info → keras_rs_nightly-0.2.2.dev202506180335.dist-info}/RECORD

@@ -5,15 +5,16 @@ keras_rs/metrics/__init__.py,sha256=Qxpf6OFooIL9TIn2l3WgOea3HFRG0hq02glPAxtMZ9c,
 keras_rs/src/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 keras_rs/src/api_export.py,sha256=RsmG-DvO-cdFeAF9W6LRzms0kvtm-Yp9BAA_d-952zI,510
 keras_rs/src/types.py,sha256=1A-oLRdX1-f2DsVZBcNl8qNsaH8pM-gnleLT9FWZWBw,1189
-keras_rs/src/version.py,sha256=ua2Hsp6lXwvYzm-yqp78ielSnR4SRHWXYVwyaSH_nj8,224
+keras_rs/src/version.py,sha256=LNFzWjtxjRhe0TX6c7WYi7sKOVa0qIKYwIc1tSfziS4,224
 keras_rs/src/layers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 keras_rs/src/layers/embedding/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-keras_rs/src/layers/embedding/base_distributed_embedding.py,sha256=dUZ4eS6ktnbnw_Z5gbyZGpQqO44Oyi7DkpNCReL66No,44347
+keras_rs/src/layers/embedding/base_distributed_embedding.py,sha256=11GicbB6m0wsHJQXISp6lcUyACVVYFLFerluUJUjDFA,44265
 keras_rs/src/layers/embedding/distributed_embedding.py,sha256=94jxUHoGK3Gs9yfV0KxFTuqPo7XFnhgCNlO2FEeiSgM,1072
 keras_rs/src/layers/embedding/distributed_embedding_config.py,sha256=AWPmZBir1shhqNP6U_jiQ9lsBhMXVikW4B5VnzLsvPg,5579
 keras_rs/src/layers/embedding/embed_reduce.py,sha256=c-MnEw1-KWs0jTf0JJ_ZBOY-9hRkiFyu989Dof3DnS8,12343
 keras_rs/src/layers/embedding/jax/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-keras_rs/src/layers/embedding/jax/config_conversion.py,sha256=kDgzab8AVYf4jd_8fsiycPA0oFnT83kSWx-TXhzy6sk,13590
+keras_rs/src/layers/embedding/jax/checkpoint_utils.py,sha256=wZ4I5WZVNg5WnrD2j7nhAXgLzDc7xMrUEkSAOx5Sz5c,3495
+keras_rs/src/layers/embedding/jax/config_conversion.py,sha256=Di1UzRwLgGHd7RuWYJMj2mCOr1u9MseFEWaYKnwD9Bs,16742
 keras_rs/src/layers/embedding/jax/distributed_embedding.py,sha256=s_V2h8smO6_Nd3lQfp6zqNi9XxXIn9wjnggSedRoE8E,35410
 keras_rs/src/layers/embedding/jax/embedding_lookup.py,sha256=HFkc0pGB9JngnCtbEJE2gDxC2K4gDdQ6GpnatSdnW6s,8205
 keras_rs/src/layers/embedding/jax/embedding_utils.py,sha256=EHrQjPLl94STLWf9g8Ew8nuwupXRq-a_QmvFlXV6G6A,20331
@@ -49,7 +50,7 @@ keras_rs/src/metrics/utils.py,sha256=fGTo8j0ykVE5Y3yQCS2orSFcHY20Uxt0NazyPsybUsw
 keras_rs/src/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 keras_rs/src/utils/doc_string_utils.py,sha256=CmqomepmaYcvpACpXEXkrJb8DMnvIgmYK-lJ53lYarY,1675
 keras_rs/src/utils/keras_utils.py,sha256=dc-NFzs3a-qmRw0vBDiMslPLfrm9yymGduLWesXPhuY,2123
-keras_rs_nightly-0.2.2.dev202506160338.dist-info/METADATA,sha256=a9xMKY-hqmkOw59DRdW4M7l6iXVR4GgRdY5qACQsXvE,5273
-keras_rs_nightly-0.2.2.dev202506160338.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-keras_rs_nightly-0.2.2.dev202506160338.dist-info/top_level.txt,sha256=pWs8X78Z0cn6lfcIb9VYOW5UeJ-TpoaO9dByzo7_FFo,9
-keras_rs_nightly-0.2.2.dev202506160338.dist-info/RECORD,,
+keras_rs_nightly-0.2.2.dev202506180335.dist-info/METADATA,sha256=yRFTgK31VJDFNIHnptUz6hmDK7NPyXdhXptXrvFN9_M,5273
+keras_rs_nightly-0.2.2.dev202506180335.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+keras_rs_nightly-0.2.2.dev202506180335.dist-info/top_level.txt,sha256=pWs8X78Z0cn6lfcIb9VYOW5UeJ-TpoaO9dByzo7_FFo,9
+keras_rs_nightly-0.2.2.dev202506180335.dist-info/RECORD,,