keras-rs-nightly 0.0.1.dev2025021903__py3-none-any.whl → 0.3.1.dev202512130338__py3-none-any.whl

keras_rs/src/layers/embedding/distributed_embedding.py

@@ -0,0 +1,33 @@
+import importlib.util
+import platform
+import sys
+
+import keras
+
+from keras_rs.src.api_export import keras_rs_export
+
+# JAX distributed embedding is only available on linux_x86_64, and only if
+# jax-tpu-embedding is installed.
+jax_tpu_embedding_spec = importlib.util.find_spec("jax_tpu_embedding")
+if (
+    keras.backend.backend() == "jax"
+    and sys.platform == "linux"
+    and platform.machine().lower() == "x86_64"
+    and jax_tpu_embedding_spec is not None
+):
+    from keras_rs.src.layers.embedding.jax.distributed_embedding import (
+        DistributedEmbedding as BackendDistributedEmbedding,
+    )
+elif keras.backend.backend() == "tensorflow":
+    from keras_rs.src.layers.embedding.tensorflow.distributed_embedding import (
+        DistributedEmbedding as BackendDistributedEmbedding,
+    )
+else:
+    from keras_rs.src.layers.embedding.base_distributed_embedding import (
+        DistributedEmbedding as BackendDistributedEmbedding,
+    )
+
+
+@keras_rs_export("keras_rs.layers.DistributedEmbedding")
+class DistributedEmbedding(BackendDistributedEmbedding):
+    pass

keras_rs/src/layers/embedding/distributed_embedding_config.py

@@ -0,0 +1,132 @@
+"""Configuration for TPU embedding layer."""
+
+import dataclasses
+from typing import Any
+
+import keras
+
+from keras_rs.src import types
+from keras_rs.src.api_export import keras_rs_export
+
+
+@keras_rs_export("keras_rs.layers.TableConfig")
+@dataclasses.dataclass(order=True)
+class TableConfig:
+    """Configuration for one embedding table.
+
+    Configures one table for use by one or more `keras_rs.layers.FeatureConfig`,
+    which in turn is used to configure a `keras_rs.layers.DistributedEmbedding`.
+
+    Attributes:
+        name: The name of the table. Must be defined.
+        vocabulary_size: Size of the table's vocabulary (number of rows).
+        embedding_dim: The embedding dimension (width) of the table.
+        initializer: The initializer for the embedding weights. If not
+            specified,  defaults to `truncated_normal_initializer` with mean
+            `0.0` and standard deviation `1 / sqrt(embedding_dim)`.
+        optimizer: The optimizer for the embedding table. Only SGD, Adagrad,
+            Adam, and FTRL are supported. Note that not all of the optimizer's
+            parameters are supported. Defaults to Adam.
+        combiner: Specifies how to reduce if there are multiple entries in a
+            single row. `mean`, `sqrtn` and `sum` are supported. `mean` is the
+            default. `sqrtn` often achieves good accuracy, in particular with
+            bag-of-words columns.
+        placement: Where to place the embedding table. `"auto"`, which is the
+            default, means that the table is placed on SparseCore if available,
+            otherwise on the default device where the rest of the model is
+            placed. A value of `"sparsecore"` means the table will be placed on
+            the SparseCore chips and an error is raised if SparseCore is not
+            available. A value of `"default_device"` means the table will be
+            placed on the default device where the rest of the model is placed,
+            even if SparseCore is available. The default device for the rest of
+            the model is the TPU's TensorCore on TPUs, otherwise the GPU or CPU.
+        max_ids_per_partition: The max number of ids per partition for the
+            table. This is an input data dependent value and is required by the
+            compiler to appropriately allocate memory.
+        max_unique_ids_per_partition: The max number of unique ids per partition
+            for the table. This is an input data dependent value and is required
+            by the compiler to appropriately allocate memory.
+    """
+
+    name: str
+    vocabulary_size: int
+    embedding_dim: int
+    initializer: str | keras.initializers.Initializer = (
+        keras.initializers.VarianceScaling(mode="fan_out")
+    )
+    optimizer: str | keras.optimizers.Optimizer = "adam"
+    combiner: str = "mean"
+    placement: str = "auto"
+    max_ids_per_partition: int = 256
+    max_unique_ids_per_partition: int = 256
+
+    def get_config(self) -> dict[str, Any]:
+        return {
+            "name": self.name,
+            "vocabulary_size": self.vocabulary_size,
+            "embedding_dim": self.embedding_dim,
+            "initializer": keras.saving.serialize_keras_object(
+                self.initializer
+            ),
+            "optimizer": keras.saving.serialize_keras_object(self.optimizer),
+            "combiner": self.combiner,
+            "placement": self.placement,
+            "max_ids_per_partition": self.max_ids_per_partition,
+            "max_unique_ids_per_partition": self.max_unique_ids_per_partition,
+        }
+
+    @classmethod
+    def from_config(cls, config: dict[str, Any]) -> "TableConfig":
+        config = config.copy()
+        config["initializer"] = keras.saving.deserialize_keras_object(
+            config["initializer"]
+        )
+        config["optimizer"] = keras.saving.deserialize_keras_object(
+            config["optimizer"]
+        )
+        return cls(**config)
+
+
+@keras_rs_export("keras_rs.layers.FeatureConfig")
+@dataclasses.dataclass(order=True)
+class FeatureConfig:
+    """Configuration for one embedding feature.
+
+    Configures one feature for `keras_rs.layers.DistributedEmbedding`. Each
+    feature uses a table configured via `keras_rs.layers.TableConfig` and
+    multiple features can share the same table.
+
+    Attributes:
+        name: The name of the feature. Must be defined.
+        table: The table in which to look up this feature.
+        input_shape: The input shape of the feature. The feature fed into the
+            layer has to match the shape. Note that for ragged dimensions in the
+            input, the dimension provided here presents the maximum value;
+            anything larger will be truncated. Also note that the first
+            dimension represents the global batch size. For example, on TPU,
+            this represents the total number of samples that are dispatched to
+            all the TPUs connected to the current host.
+        output_shape: The output shape of the feature activation. What is
+            returned by the embedding layer has to match this shape.
+    """
+
+    name: str
+    table: TableConfig
+    input_shape: types.Shape
+    output_shape: types.Shape
+
+    def get_config(self) -> dict[str, Any]:
+        return {
+            "name": self.name,
+            "table": self.table.get_config(),
+            "input_shape": self.input_shape,
+            "output_shape": self.output_shape,
+        }
+
+    @classmethod
+    def from_config(cls, config: dict[str, Any]) -> "FeatureConfig":
+        config = config.copy()
+        # Note: the handling of shared tables during serialization is done in
+        # `DistributedEmbedding.from_config()`.
+        config["table"] = TableConfig.from_config(config["table"])
+        return cls(**config)

keras_rs/src/layers/embedding/embed_reduce.py

@@ -0,0 +1,309 @@
+from typing import Any
+
+import keras
+from keras import ops
+
+from keras_rs.src import types
+from keras_rs.src.api_export import keras_rs_export
+from keras_rs.src.utils.keras_utils import check_shapes_compatible
+
+SUPPORTED_COMBINERS = ("mean", "sum", "sqrtn")
+
+
+def _is_supported_sparse(x: types.Tensor) -> bool:
+    """Determines if the input is a supported sparse tensor.
+
+    NOTE: Currently only works for the TensorFlow and JAX backends.
+
+    Args:
+      x: Input tensor to check for sparsity.
+
+    Returns:
+      True if `x` is a supported sparse tensor.
+    """
+    if keras.backend.backend() == "tensorflow":
+        import tensorflow as tf
+
+        return isinstance(x, tf.SparseTensor)
+    elif keras.backend.backend() == "jax":
+        from jax.experimental import sparse as jax_sparse
+
+        return isinstance(x, jax_sparse.BCOO) or isinstance(x, jax_sparse.BCSR)
+
+    return False
+
+
+def _sparse_ones_like(
+    x: types.Tensor, dtype: types.DType | None = None
+) -> types.Tensor:
+    """Creates a tensor of ones with the same sparsity as the input.
+
+    This differs from `keras.ops.ones_like`, which would create a dense
+    tensor of ones.
+
+    Args:
+        x: Input sparse tensor.
+        dtype: Optional dtype for the output tensor values.
+
+    Returns:
+        Sparse tensor of ones.
+
+    Raises:
+        ValueError for unsupported sparse input type and backend.
+    """
+    dtype = dtype or x.dtype
+    if keras.backend.backend() == "tensorflow":
+        import tensorflow as tf
+
+        # Ensure shape is copied exactly for compatibility in graph mode.
+        x_shape = x.shape
+        y = tf.SparseTensor(
+            x.indices, tf.ones_like(x.values, dtype=dtype), x.dense_shape
+        )
+        y.set_shape(x_shape)
+        return y
+    elif keras.backend.backend() == "jax":
+        import jax.numpy as jnp
+        from jax.experimental import sparse as jax_sparse
+
+        if isinstance(x, jax_sparse.BCOO):
+            return jax_sparse.BCOO(
+                (jnp.ones_like(x.data, dtype=dtype), x.indices),
+                shape=x.shape,
+                indices_sorted=x.indices_sorted,
+                unique_indices=x.unique_indices,
+            )
+        elif isinstance(x, jax_sparse.BCSR):
+            return jax_sparse.BCSR(
+                (jnp.ones_like(x.data, dtype=dtype), x.indices, x.indptr),
+                shape=x.shape,
+                indices_sorted=x.indices_sorted,
+                unique_indices=x.unique_indices,
+            )
+
+    raise ValueError(
+        f"Unsupported sparse input type '{x.__class__.__name__}' for backend "
+        f"{keras.backend.backend()}."
+    )
+
+
+@keras_rs_export("keras_rs.layers.EmbedReduce")
+class EmbedReduce(keras.layers.Embedding):
+    """An embedding layer that reduces with a combiner.
+
+    This layer embeds inputs and then applies a reduction to combine a set of
+    embeddings into a single embedding. This is typically used to embed a
+    sequence of items as a single embedding.
+
+    If the inputs passed to `__call__` are 1D, no reduction is applied. If the
+    inputs are 2D, dimension 1 is reduced using the combiner so that the result
+    is of shape `(batch_size, output_dim`). Inputs of rank 3 and higher are not
+    allowed. Weights can optionally be passed to the `__call__` method to
+    apply weights to different samples before reduction.
+
+    This layer supports sparse inputs and ragged inputs with backends that
+    support them. The output after reduction is dense. For ragged inputs, the
+    ragged dimension must be 1 as it is the dimension that is reduced.
+
+    Args:
+        input_dim: Integer. Size of the vocabulary, maximum integer index + 1.
+        output_dim: Integer. Dimension of the dense embedding.
+        embeddings_initializer: Initializer for the `embeddings` matrix (see
+            `keras.initializers`).
+        embeddings_regularizer: Regularizer function applied to the `embeddings`
+            matrix (see `keras.regularizers`).
+        embeddings_constraint: Constraint function applied to the `embeddings`
+            matrix (see `keras.constraints`).
+        mask_zero: Boolean, whether or not the input value 0 is a special
+            "padding" value that should be masked out. This is useful when using
+            recurrent layers which may take variable length input. If this is
+            `True`, then all subsequent layers in the model need to support
+            masking or an exception will be raised. If `mask_zero` is set to
+            `True`, as a consequence, index 0 cannot be used in the vocabulary
+            (`input_dim` should equal size of vocabulary + 1).
+        weights: Optional floating-point matrix of size
+            `(input_dim, output_dim)`. The initial embeddings values to use.
+        combiner: Specifies how to reduce if there are multiple entries in a
+            single row. Currently `mean`, `sqrtn` and `sum` are supported.
+            `mean` is the default. `sqrtn` often achieves good accuracy, in
+            particular with bag-of-words columns.
+        **kwargs: Additional keyword arguments passed to `Embedding`.
+    """
+
+    def __init__(
+        self,
+        input_dim: int,
+        output_dim: int,
+        embeddings_initializer: types.InitializerLike = "uniform",
+        embeddings_regularizer: types.RegularizerLike | None = None,
+        embeddings_constraint: types.ConstraintLike | None = None,
+        mask_zero: bool = False,
+        weights: types.Tensor = None,
+        combiner: str = "mean",
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(
+            input_dim,
+            output_dim,
+            embeddings_initializer=embeddings_initializer,
+            embeddings_regularizer=embeddings_regularizer,
+            embeddings_constraint=embeddings_constraint,
+            mask_zero=mask_zero,
+            weights=weights,
+            **kwargs,
+        )
+        if combiner not in SUPPORTED_COMBINERS:
+            raise ValueError(
+                f"Invalid `combiner`: '{combiner}', "
+                f"use one of {', '.join(SUPPORTED_COMBINERS)}."
+            )
+        self.combiner = combiner
+
+    def call(
+        self,
+        inputs: types.Tensor,
+        weights: types.Tensor | None = None,
+    ) -> types.Tensor:
+        """Apply embedding and reduction.
+
+        Args:
+            inputs: 1D tensor to embed or 2D tensor to embed and reduce.
+            weights: Optional tensor of weights to apply before reduction, which
+               can be 1D or 2D and must match for the first dimension of
+               `inputs` (1D case) or match the shape of `inputs` (2D case).
+
+        Returns:
+            A dense 2D tensor of shape `(batch_size, output_dim)`.
+        """
+        x = super().call(inputs)
+        unreduced_rank = len(x.shape)
+
+        # Check that weights has a compatible shape.
+        if weights is not None:
+            weights_rank = len(weights.shape)
+            if weights_rank > unreduced_rank or not check_shapes_compatible(
+                x.shape[0:weights_rank], weights.shape
+            ):
+                raise ValueError(
+                    f"The shape of `weights`: {weights.shape} is not compatible"
+                    f" with the shape of `inputs` after embedding: {x.shape}."
+                )
+
+        dtype = (
+            x.dtype
+            if weights is None
+            else keras.backend.result_type(x.dtype, weights.dtype)
+        )
+
+        # When `weights` is `None`:
+        # - For ragged inputs, after embedding, we get a ragged result that has
+        #   a ragged dimension of 1, but when we do the "mean" or "sqrtn", we
+        #   need to divide by the number of items in each row. However, there is
+        #   no explicit cross backend API to get the row length. `ones_like`
+        #   gives us a ragged tensor that is ragged in the same way as the
+        #   inputs. When we do `ops.sum(weights, axis=-2)`, it gives us the
+        #   number of items per row.
+        # - For sparse inputs, after embedding, we get a dense tensor, not a
+        #   sparse tensor. What it does for missing values is use embedding 0.
+        #   These are bogus embedding and should be ignored. `ones_like` gives
+        #   us a sparse tensor with the exact same missing values. Later, when
+        #   we do `x = ops.multiply(x, weights)`, which masks the bogus values
+        #   (note that `weights` has been densified beforehand). Additionally,
+        #   when we do `ops.sum(weights, axis=-2)`, it gives us the number of
+        #   items per row.
+        #
+        # When `unreduced_rank <= 2`, this means that the inputs where 1D and
+        # dense, there is only one embedding per row, so there is no real
+        # reduction is going on.
+        # - For mean: result = weights * x / weights = x we don't need `weights`
+        # - For sqrtn: result = weights * x / sqrt(square(weights)) = x we don't
+        #   needs `weights`
+        # - For sum however: `result = weights * x` we do need `weights`.
+        # So for mean and sqrtn we don't need the weights, we use ones instead.
+        # This is to avoid divisions by zero and improve the precision.
+        if weights is None or (unreduced_rank <= 2 and self.combiner != "sum"):
+            # Discard the weights if there were some and create a mask for
+            # ragged and sparse tensors to mask the result correctly (sparse
+            # only) and the apply the reduction correctly (ragged and sparse).
+            if _is_supported_sparse(inputs):
+                weights = _sparse_ones_like(inputs, dtype=dtype)
+            else:
+                weights = ops.ones_like(inputs, dtype=dtype)
+
+        else:
+            weights = ops.cast(weights, dtype)
+
+        # When looking up using sparse indices, the result is dense but contains
+        # values that should be ignored as all missing values use index 0. We
+        # use `weights` as a mask, but it needs to be densified as
+        # `expand_dims` and broadcasting a sparse tensor does not produce the
+        # expected result.
+        weights = ops.convert_to_tensor(weights, sparse=False)
+
+        # Make weights and the unreduced embeddings have the same rank.
+        weights_rank = len(weights.shape)
+        if weights_rank < unreduced_rank:
+            weights = ops.expand_dims(
+                weights, axis=tuple(range(weights_rank, unreduced_rank))
+            )
+
+        # Note that `x` and `weights` are:
+        # - ragged if `inputs` was ragged and `weights` was ragged or None
+        # - dense otherwise (even if `inputs` and `weights` were sparse).
+        x = ops.multiply(x, weights)
+
+        if unreduced_rank <= 2:
+            # No reduction is applied.
+            return x
+
+        # After this reduction, `x` is always dense as we reduce the ragged
+        # dimension in the ragged case.
+        x = ops.sum(x, axis=-2)
+
+        # Apply the right divisor for the combiner.
+        # Where we use `weights` in the divisor, we use
+        # `ops.sum(weights, axis=-2)` which always makes it dense as we reduce
+        # the ragged dimension in the ragged case.
+        if self.combiner == "mean":
+            return ops.divide_no_nan(x, ops.sum(weights, axis=-2))
+        elif self.combiner == "sum":
+            return x
+        elif self.combiner == "sqrtn":
+            return ops.divide_no_nan(
+                x, ops.sqrt(ops.sum(ops.square(weights), axis=-2))
+            )
+
+    def compute_output_shape(
+        self,
+        input_shape: types.Shape,
+        weights_shape: types.Shape | None = None,
+    ) -> types.Shape:
+        del weights_shape
+
+        if len(input_shape) <= 1:
+            # No reduction
+            return (*input_shape, self.output_dim)
+        else:
+            # Reduce last dimension
+            return (*input_shape[0:-1], self.output_dim)
+
+    def compute_output_spec(
+        self,
+        inputs: keras.KerasTensor,
+        weights: keras.KerasTensor | None = None,
+    ) -> keras.KerasTensor:
+        del weights
+
+        output_shape = self.compute_output_shape(inputs.shape)
+        return keras.KerasTensor(output_shape, dtype=self.compute_dtype)
+
+    def get_config(self) -> dict[str, Any]:
+        config: dict[str, Any] = super().get_config()
+
+        config.update(
+            {
+                "combiner": self.combiner,
+            }
+        )
+
+        return config

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)