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

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

@@ -0,0 +1,276 @@
+"""Defines a differentiable `embedding_lookup` function.
+
+Implementation details for use in JAX models.
+"""
+
+import functools
+from typing import Any, Mapping, TypeAlias
+
+import jax
+import numpy as np
+from jax.experimental import layout as jax_layout
+from jax_tpu_embedding.sparsecore.lib.nn import embedding
+from jax_tpu_embedding.sparsecore.lib.nn import embedding_spec
+from jax_tpu_embedding.sparsecore.utils import utils as jte_utils
+
+from keras_rs.src.layers.embedding.jax import embedding_utils
+from keras_rs.src.types import Nested
+
+if jax.__version_info__ >= (0, 8, 0):
+    from jax import shard_map
+else:
+    from jax.experimental.shard_map import shard_map as exp_shard_map
+
+    def shard_map(  # type: ignore[misc]
+        f: Any = None,
+        /,
+        *,
+        out_specs: Any,
+        in_specs: Any,
+        mesh: Any = None,
+        check_vma: bool = True,
+    ) -> Any:
+        return exp_shard_map(
+            f,
+            mesh=mesh,
+            in_specs=in_specs,
+            out_specs=out_specs,
+            check_rep=check_vma,
+        )  # type: ignore[no-untyped-call]
+
+
+ShardedCooMatrix = embedding_utils.ShardedCooMatrix
+ArrayLike: TypeAlias = jax.Array | np.ndarray[Any, Any]
+JaxLayout: TypeAlias = jax.sharding.NamedSharding | jax_layout.Format
+
+
+class EmbeddingLookupConfiguration:
+    """Feature, mesh and sharding configrmation for lookups."""
+
+    mesh: jax.sharding.Mesh
+    feature_specs: embedding.Nested[embedding_spec.FeatureSpec]
+    table_sharding_strategy: str
+    num_sc_per_device: int
+    samples_partition: jax.sharding.PartitionSpec
+    table_partition: jax.sharding.PartitionSpec
+    table_layout: JaxLayout
+
+    def __init__(
+        self,
+        feature_specs: embedding.Nested[embedding_spec.FeatureSpec],
+        mesh: jax.sharding.Mesh | None = None,
+        table_sharding_strategy: str = "MOD",
+        num_sc_per_device: int | None = None,
+        sharding_axis: str = "sparsecore_sharding",
+        samples_partition: jax.sharding.PartitionSpec | None = None,
+        samples_layout: JaxLayout | None = None,
+        table_partition: jax.sharding.PartitionSpec | None = None,
+        table_layout: JaxLayout | None = None,
+    ):
+        self.mesh = mesh or jax.sharding.Mesh(jax.devices(), sharding_axis)
+        self.feature_specs = feature_specs
+        self.table_sharding_strategy = table_sharding_strategy
+        self.num_sc_per_device = (
+            num_sc_per_device
+            if num_sc_per_device is not None
+            else jte_utils.num_sparsecores_per_device()
+        )
+        self.samples_partition = (
+            samples_partition
+            or jax.sharding.PartitionSpec(
+                sharding_axis  # type: ignore[no-untyped-call]
+            )
+        )
+        self.samples_layout = samples_layout or jax.sharding.NamedSharding(
+            self.mesh, self.samples_partition
+        )
+        self.table_partition = table_partition or jax.sharding.PartitionSpec(
+            sharding_axis,
+            None,  # type: ignore[no-untyped-call]
+        )
+        self.table_layout = table_layout or jax.sharding.NamedSharding(
+            self.mesh, self.table_partition
+        )
+
+
+# Embedding lookup function with custom gradient.
+@functools.partial(jax.custom_vjp, nondiff_argnums=(0,))
+def embedding_lookup(
+    config: EmbeddingLookupConfiguration,
+    lookups: Mapping[str, ShardedCooMatrix],
+    tables: Nested[jax.Array],
+    step: jax.Array | None = None,
+) -> Nested[jax.Array]:
+    """Embedding lookup function with custom gradient.
+
+    Args:
+      config: Embedding lookup configuration.
+      lookups: Embedding lookup stacked/sharded inputs.
+      tables: Embedding lookup stacked/sharded tables.
+      step: Current training step number.
+    """
+    del step  # Only used in backward pass.
+
+    # Decompose COO matrices.
+    row_pointers_raw = {}
+    embedding_ids_raw = {}
+    sample_ids_raw = {}
+    gains_raw = {}
+    for table_name, coo in lookups.items():
+        row_pointers_raw[table_name] = coo.shard_ends
+        embedding_ids_raw[table_name] = coo.col_ids
+        sample_ids_raw[table_name] = coo.row_ids
+        gains_raw[table_name] = coo.values
+
+    sparse_dense_matmul_input = embedding.SparseDenseMatmulInput(
+        row_pointers_raw,
+        embedding_ids_raw,
+        sample_ids_raw,
+        gains_raw,
+    )
+
+    pd = config.samples_partition
+    pt = config.table_partition
+    sharded_matmul = jax.jit(
+        shard_map(
+            functools.partial(
+                embedding.tpu_sparse_dense_matmul,
+                global_device_count=config.mesh.shape[pd[0]],
+                feature_specs=config.feature_specs,
+                sharding_strategy=config.table_sharding_strategy,
+            ),
+            mesh=config.mesh,
+            in_specs=(pd, pt),
+            out_specs=pd,
+            check_vma=False,
+        ),
+    )
+
+    activations: Nested[jax.Array] = sharded_matmul(
+        sparse_dense_matmul_input,
+        tables,
+    )
+
+    return activations
+
+
+def embedding_lookup_fwd(
+    config: EmbeddingLookupConfiguration,
+    lookups: Mapping[str, ShardedCooMatrix],
+    table: Nested[jax.Array],
+    step: jax.Array | None = None,
+) -> tuple[
+    Nested[jax.Array],
+    tuple[Nested[ShardedCooMatrix], Nested[jax.Array], jax.Array | None],
+]:
+    """Forward pass for embedding lookup."""
+    return embedding_lookup(config, lookups, table, step), (
+        lookups,
+        table,
+        step,
+    )
+
+
+def embedding_lookup_bwd(
+    config: EmbeddingLookupConfiguration,
+    res: tuple[
+        Mapping[str, ShardedCooMatrix],  # Lookups.
+        Mapping[str, Nested[jax.Array]],  # Tables.
+        jax.Array | None,  # Step.
+    ],
+    gradients: Nested[jax.Array],
+) -> tuple[None, Nested[jax.Array], jax.Array | None]:
+    """Backward pass for embedding lookup.
+
+    Args:
+      config: Embedding lookup configuration.
+      res: Tuple of embedding lookup (inputs, tables, step).
+      gradients: Embedding lookup gradients.
+
+    Returns:
+      A tuple of gradients (None, table_grads, step + 1).
+    """
+    lookups, tables, step = res
+
+    # Decompose COO matrices.
+    row_pointers_raw = {}
+    embedding_ids_raw = {}
+    sample_ids_raw = {}
+    gains_raw = {}
+    for table_name, coo in lookups.items():
+        row_pointers_raw[table_name] = coo.shard_ends
+        embedding_ids_raw[table_name] = coo.col_ids
+        sample_ids_raw[table_name] = coo.row_ids
+        gains_raw[table_name] = coo.values
+
+    sparse_dense_matmul_input = embedding.SparseDenseMatmulInput(
+        row_pointers_raw,
+        embedding_ids_raw,
+        sample_ids_raw,
+        gains_raw,
+    )
+
+    pt = config.table_partition
+    pd = config.samples_partition
+    # Replicate step count.
+    preplicate = jax.sharding.PartitionSpec()  # type: ignore[no-untyped-call]
+
+    def grad_func(
+        gradients: Nested[jax.Array],
+        sparse_input: embedding.SparseDenseMatmulInput,
+        tables: Mapping[str, embedding.EmbeddingVariables],
+        step: jax.Array | None,
+    ) -> Mapping[str, embedding.EmbeddingVariables]:
+        output: Mapping[str, embedding.EmbeddingVariables] = (
+            embedding.tpu_sparse_dense_matmul_grad(
+                gradients,
+                sparse_input,
+                tables,
+                feature_specs=config.feature_specs,
+                sharding_strategy=config.table_sharding_strategy,
+                step=step,
+            )
+        )
+        return output
+
+    # activation_layout = jax.sharding.NamedSharding(config.mesh, pd)
+    # step_layout = jax.sharding.NamedSharding(config.mesh, preplicate)
+    sharded_matmul_grad = jax.jit(
+        shard_map(
+            grad_func,
+            mesh=config.mesh,
+            in_specs=(pd, pd, pt, preplicate),
+            out_specs=pt,
+            check_vma=False,
+        ),
+        #   in_shardings=(
+        #       activation_layout,
+        #       config.samples_layout,
+        #       config.table_layout,
+        #       step_layout,
+        #   ),
+        # out_shardings=(config.table_layout),
+    )
+
+    table_grads = sharded_matmul_grad(
+        gradients,
+        sparse_dense_matmul_input,
+        tables,
+        step,
+    )
+
+    # tpu_sparse_dense_matmul_grad returns a general Mapping (usually a dict).
+    # It may not be the same type as the embedding table (e.g. FrozenDict).
+    # Here we use flatten / unflatten to ensure the types are the same.
+    table_grads = jax.tree.unflatten(
+        jax.tree.structure(tables), jax.tree.leaves(table_grads)
+    )
+
+    return (
+        None,
+        table_grads,
+        step + 1 if step is not None else None,  # Incremented step count.
+    )
+
+
+embedding_lookup.defvjp(embedding_lookup_fwd, embedding_lookup_bwd)

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

@@ -0,0 +1,217 @@
+"""Utility functions for manipulating JAX embedding tables and inputs."""
+
+import collections
+from typing import Any, Mapping, NamedTuple, Sequence, TypeAlias, TypeVar
+
+import jax
+import numpy as np
+from jax_tpu_embedding.sparsecore.lib.nn import embedding
+from jax_tpu_embedding.sparsecore.lib.nn import table_stacking
+from jax_tpu_embedding.sparsecore.lib.nn.embedding_spec import FeatureSpec
+
+from keras_rs.src.types import Nested
+
+T = TypeVar("T")
+
+# Any to support tf.Ragged without needing an explicit TF dependency.
+ArrayLike: TypeAlias = jax.Array | np.ndarray | Any  # type: ignore
+Shape: TypeAlias = tuple[int, ...]
+
+
+class FeatureSamples(NamedTuple):
+    tokens: ArrayLike
+    weights: ArrayLike | None
+
+
+class ShardedCooMatrix(NamedTuple):
+    shard_starts: ArrayLike
+    shard_ends: ArrayLike
+    col_ids: ArrayLike
+    row_ids: ArrayLike
+    values: ArrayLike
+
+
+def convert_to_numpy(
+    ragged_or_dense: np.ndarray[Any, Any] | Sequence[Sequence[Any]] | Any,
+    dtype: Any,
+) -> np.ndarray[Any, Any]:
+    """Converts a ragged or dense list of inputs to a ragged/dense numpy array.
+
+    The output is adjusted to be 2D.
+
+    Args:
+        ragged_or_dense: Input that is either already a numpy array, or nested
+            sequence.
+        dtype: Numpy dtype of output array.
+
+    Returns:
+        Corresponding numpy array.
+    """
+    if hasattr(ragged_or_dense, "numpy"):
+        # Support tf.RaggedTensor and other TF input dtypes.
+        if callable(getattr(ragged_or_dense, "numpy")):
+            ragged_or_dense = ragged_or_dense.numpy()
+
+    if isinstance(ragged_or_dense, jax.Array):
+        ragged_or_dense = np.asarray(ragged_or_dense)
+
+    if isinstance(ragged_or_dense, np.ndarray):
+        # Convert 1D to 2D.
+        if ragged_or_dense.dtype != np.ndarray and ragged_or_dense.ndim == 1:
+            return ragged_or_dense.reshape(-1, 1).astype(dtype)
+
+        # If dense, return converted dense type.
+        if ragged_or_dense.dtype != np.ndarray:
+            return ragged_or_dense.astype(dtype)
+
+        # Ragged numpy array.
+        return ragged_or_dense
+
+    # Handle 1D sequence input.
+    if not isinstance(ragged_or_dense[0], collections.abc.Sequence):
+        return np.asarray(ragged_or_dense, dtype=dtype).reshape(-1, 1)
+
+    # Assemble elements into an nd-array.
+    counts = [len(vals) for vals in ragged_or_dense]
+    if all([count == counts[0] for count in counts]):
+        # Dense input.
+        return np.asarray(ragged_or_dense, dtype=dtype)
+    else:
+        # Ragged input, convert to ragged numpy arrays.
+        return np.array(
+            [np.array(row, dtype=dtype) for row in ragged_or_dense],
+            dtype=np.ndarray,
+        )
+
+
+def create_feature_samples(
+    feature_structure: Nested[T],
+    feature_ids: Nested[ArrayLike | Sequence[int] | Sequence[Sequence[int]]],
+    feature_weights: None
+    | (Nested[ArrayLike | Sequence[float] | Sequence[Sequence[float]]]),
+) -> Nested[FeatureSamples]:
+    """Constructs a collection of sample tuples from provided IDs and weights.
+
+    Args:
+        feature_structure: The nested structure of the inputs (typically
+            `FeatureSpec`s).
+        feature_ids: The feature IDs to use for the samples.
+        feature_weights: The feature weights to use for the samples.  Defaults
+            to ones if not provided.
+
+    Returns:
+        A nested collection of `FeatureSamples` corresponding to the input IDs
+        and weights, for use in embedding lookups.
+    """
+    # Create numpy arrays from inputs.
+    feature_ids = jax.tree.map(
+        lambda _, ids: convert_to_numpy(ids, np.int32),
+        feature_structure,
+        feature_ids,
+    )
+
+    if feature_weights is None:
+        return jax.tree.map(  # type: ignore[no-any-return]
+            lambda _, ids: FeatureSamples(ids, None),
+            feature_structure,
+            feature_ids,
+        )
+
+    feature_weights = jax.tree.map(
+        lambda _, wgts: convert_to_numpy(wgts, np.float32),
+        feature_structure,
+        feature_weights,
+    )
+
+    # Assemble.
+    def _create_feature_samples(
+        sample_ids: np.ndarray[Any, Any],
+        sample_weights: np.ndarray[Any, Any],
+    ) -> FeatureSamples:
+        return FeatureSamples(sample_ids, sample_weights)
+
+    output: Nested[FeatureSamples] = jax.tree.map(
+        lambda _, sample_ids, sample_weights: _create_feature_samples(
+            sample_ids, sample_weights
+        ),
+        feature_structure,
+        feature_ids,
+        feature_weights,
+    )
+    return output
+
+
+def stack_and_shard_samples(
+    feature_specs: Nested[FeatureSpec],
+    feature_samples: Nested[FeatureSamples],
+    local_device_count: int,
+    global_device_count: int,
+    num_sc_per_device: int,
+    static_buffer_size: int | Mapping[str, int] | None = None,
+) -> tuple[dict[str, ShardedCooMatrix], embedding.SparseDenseMatmulInputStats]:
+    """Prepares input samples for use in embedding lookups.
+
+    Args:
+        feature_specs: Nested collection of feature specifications.
+        feature_samples: Nested collection of feature samples.
+        local_device_count: Number of local JAX devices.
+        global_device_count: Number of global JAX devices.
+        num_sc_per_device: Number of sparsecores per device.
+        static_buffer_size: The static buffer size to use for the samples.
+          Defaults to None, in which case an upper-bound for the buffer size
+          will be automatically determined.
+
+    Returns:
+        The preprocessed inputs, and statistics useful for updating FeatureSpecs
+        based on the provided input data.
+    """
+    del static_buffer_size  # Currently ignored.
+    flat_feature_specs, _ = jax.tree.flatten(feature_specs)
+
+    feature_tokens = []
+    collected_weights = []
+
+    def collect_tokens_and_weights(
+        feature_spec: FeatureSpec, samples: FeatureSamples
+    ) -> None:
+        del feature_spec
+        feature_tokens.append(samples.tokens)
+        collected_weights.append(samples.weights)
+
+    jax.tree.map(collect_tokens_and_weights, feature_specs, feature_samples)
+
+    feature_weights = (
+        None if all(w is None for w in collected_weights) else collected_weights
+    )
+
+    preprocessed_inputs, stats = embedding.preprocess_sparse_dense_matmul_input(
+        feature_tokens,
+        feature_weights,
+        flat_feature_specs,
+        local_device_count=local_device_count,
+        global_device_count=global_device_count,
+        num_sc_per_device=num_sc_per_device,
+        sharding_strategy="MOD",
+        has_leading_dimension=False,
+        allow_id_dropping=True,
+    )
+
+    out: dict[str, ShardedCooMatrix] = {}
+    tables_names = preprocessed_inputs.lhs_row_pointers.keys()
+    for table_name in tables_names:
+        shard_ends = preprocessed_inputs.lhs_row_pointers[table_name]
+        shard_starts = np.concatenate(
+            [
+                np.asarray([0]),
+                table_stacking._next_largest_multiple(shard_ends[:-1], 8),
+            ]
+        )
+        out[table_name] = ShardedCooMatrix(
+            shard_starts=shard_starts,
+            shard_ends=shard_ends,
+            col_ids=preprocessed_inputs.lhs_embedding_ids[table_name],
+            row_ids=preprocessed_inputs.lhs_sample_ids[table_name],
+            values=preprocessed_inputs.lhs_gains[table_name],
+        )
+
+    return out, stats