brainstate 0.1.10__py2.py3-none-any.whl → 0.2.1__py2.py3-none-any.whl

brainstate/nn/_embedding.py

@@ -1,58 +1,408 @@
-# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-
-from typing import Optional, Callable, Union
-
-from brainstate import init
-from brainstate._state import ParamState
-from brainstate.typing import ArrayLike
-from ._module import Module
-
-__all__ = [
-    'Embedding',
-]
-
-
-class Embedding(Module):
-    r"""
-    A simple lookup table that stores embeddings of a fixed size.
-
-    Args:
-      num_embeddings: Size of embedding dictionary. Must be non-negative.
-      embedding_size: Size of each embedding vector. Must be non-negative.
-      embedding_init: The initializer for the embedding lookup table, of shape `(num_embeddings, embedding_size)`.
-    """
-
-    def __init__(
-        self,
-        num_embeddings: int,
-        embedding_size: int,
-        embedding_init: Union[Callable, ArrayLike] = init.LecunUniform(),
-        name: Optional[str] = None,
-    ):
-        super().__init__(name=name)
-        if num_embeddings < 0:
-            raise ValueError("num_embeddings must not be negative.")
-        if embedding_size < 0:
-            raise ValueError("embedding_size must not be negative.")
-        self.num_embeddings = num_embeddings
-        self.embedding_size = embedding_size
-        self.out_size = (embedding_size,)
-
-        weight = init.param(embedding_init, (self.num_embeddings, self.embedding_size))
-        self.weight = ParamState(weight)
-
-    def update(self, indices: ArrayLike):
-        return self.weight.value[indices]
+# Copyright 2024 BrainX Ecosystem Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from functools import lru_cache
+from typing import Callable, Optional, Tuple, Union
+
+import numpy as np
+import jax
+import jax.numpy as jnp
+import jax.tree_util as jtu
+from jax import core as jax_core
+
+from brainstate._state import ParamState
+from brainstate.typing import ArrayLike, Size
+from . import init as init
+from ._module import Module
+
+__all__ = [
+    'Embedding',
+]
+
+
+def _normalize_embedding_size(size: Size) -> Tuple[int, ...]:
+    """Convert ``Size`` specifications to a validated tuple of integers."""
+    size_array = np.asarray(size)
+    if size_array.size == 0:
+        raise ValueError('embedding_size must contain at least one dimension.')
+    flat = size_array.reshape(-1)
+    normalized = tuple(int(dim) for dim in flat)
+    if any(dim < 0 for dim in normalized):
+        raise ValueError('embedding_size must not contain negative values.')
+    return normalized
+
+
+@lru_cache(maxsize=None)
+def _embedding_lookup_fn(
+    padding_idx: Optional[int],
+    scale_grad_by_freq: bool,
+):
+    """Return a lookup function with a custom VJP implementing embedding semantics."""
+
+    @jax.custom_vjp
+    def _lookup(weight: jax.Array, indices: jax.Array) -> jax.Array:
+        indices = jnp.asarray(indices)
+        return weight[indices]
+
+    def _lookup_fwd(weight: jax.Array, indices: jax.Array):
+        indices = jnp.asarray(indices)
+        return weight[indices], (indices, weight.shape)
+
+    def _lookup_bwd(residual, grad_output: jax.Array):
+        indices, weight_shape = residual
+        grad_output = jnp.asarray(grad_output)
+        flat_idx = jnp.ravel(indices)
+        if flat_idx.size == 0:
+            return jnp.zeros(weight_shape, dtype=grad_output.dtype), None
+
+        flat_idx = jnp.asarray(flat_idx, dtype=jnp.int32)
+        grad_flat = grad_output.reshape((flat_idx.shape[0],) + weight_shape[1:])
+
+        if scale_grad_by_freq:
+            counts = jnp.bincount(flat_idx, length=weight_shape[0])
+            counts = counts.astype(grad_flat.dtype)
+            counts = jnp.where(counts == 0, 1.0, counts)
+            scale = counts[flat_idx]
+            grad_flat = grad_flat / scale.reshape((flat_idx.shape[0],) + (1,) * (grad_flat.ndim - 1))
+
+        if padding_idx is not None:
+            pad_value = jnp.asarray(padding_idx, dtype=flat_idx.dtype)
+            mask = flat_idx != pad_value
+            broadcast_shape = (flat_idx.shape[0],) + (1,) * (grad_flat.ndim - 1)
+            grad_flat = grad_flat * mask.reshape(broadcast_shape).astype(grad_flat.dtype)
+
+        grad_weight = jnp.zeros(weight_shape, dtype=grad_output.dtype)
+        grad_weight = grad_weight.at[flat_idx].add(grad_flat)
+        return grad_weight, None
+
+    _lookup.defvjp(_lookup_fwd, _lookup_bwd)
+    return _lookup
+
+
+def _contains_tracer(tree) -> bool:
+    """Return True if the pytree contains any JAX tracer values."""
+    return any(isinstance(leaf, jax_core.Tracer) for leaf in jtu.tree_leaves(tree))
+
+
+
+class Embedding(Module):
+    r"""
+    A simple lookup table that stores embeddings of a fixed size.
+
+    This module is commonly used to store word embeddings and retrieve them using indices.
+    The input to the module is a list of indices, and the output is the corresponding
+    word embeddings.
+
+    Parameters
+    ----------
+    num_embeddings : int
+        Size of embedding dictionary. Must be non-negative.
+    embedding_size : Size
+        Size of each embedding vector. Can be an int or a sequence of ints, and must
+        contain only non-negative values.
+    embedding_init : Callable or ArrayLike, optional
+        The initializer for the embedding lookup table, of shape
+        ``(num_embeddings, embedding_size)``. Default is ``LecunUniform()``.
+    padding_idx : int, optional
+        If specified, the entries at ``padding_idx`` do not contribute to the gradient;
+        therefore, the embedding vector at ``padding_idx`` is not updated during training,
+        i.e., it remains as a fixed "pad". For a newly constructed Embedding, the embedding
+        vector at ``padding_idx`` will default to all zeros. Default is ``None``.
+    max_norm : float, optional
+        If given, each embedding vector with norm larger than ``max_norm`` is renormalized
+        to have norm ``max_norm``. Default is ``None``.
+    norm_type : float, optional
+        The p of the p-norm to compute for the ``max_norm`` option. Default is ``2.0``.
+    scale_grad_by_freq : bool, optional
+        If given, this scales gradients by the inverse frequency of the words in
+        the mini-batch. Default is ``False``.
+    name : str, optional
+        The name of the module.
+    param_type : type, optional
+        The parameter state type to use. Default is ``ParamState``.
+
+    Attributes
+    ----------
+    num_embeddings : int
+        Size of the embedding dictionary.
+    embedding_size : tuple[int, ...]
+        Size of each embedding vector.
+    out_size : tuple[int, ...]
+        Output size, same as ``embedding_size``.
+    weight : ParamState
+        The learnable weights of the module of shape
+        ``(num_embeddings, *embedding_size)``.
+    padding_idx : int or None
+        Index of the padding token.
+    max_norm : float or None
+        Maximum norm for embedding vectors.
+    norm_type : float
+        Type of p-norm to compute for max_norm.
+    scale_grad_by_freq : bool
+        Whether to scale gradients by frequency.
+    freeze : bool
+        Whether the embedding weights are frozen.
+
+    Examples
+    --------
+    Create an embedding layer with 10 words and 3-dimensional embeddings:
+
+    .. code-block:: python
+
+        >>> import brainstate as bst
+        >>> embedding = bst.nn.Embedding(num_embeddings=10, embedding_size=3)
+        >>> embedding.weight.value.shape
+        (10, 3)
+
+    Retrieve embeddings for specific indices:
+
+    .. code-block:: python
+
+        >>> import jax.numpy as jnp
+        >>> indices = jnp.array([1, 3, 5])
+        >>> output = embedding(indices)
+        >>> output.shape
+        (3, 3)
+
+    Use with a batch of sequences:
+
+    .. code-block:: python
+
+        >>> # Batch of 2 sequences, each with 4 tokens
+        >>> batch_indices = jnp.array([[1, 2, 3, 4],
+        ...                            [5, 6, 7, 8]])
+        >>> output = embedding(batch_indices)
+        >>> output.shape
+        (2, 4, 3)
+
+    Use padding_idx to keep padding embeddings fixed:
+
+    .. code-block:: python
+
+        >>> embedding = bst.nn.Embedding(num_embeddings=10, embedding_size=3, padding_idx=0)
+        >>> # The embedding at index 0 will remain zeros and not be updated during training
+        >>> indices = jnp.array([0, 2, 0, 5])
+        >>> output = embedding(indices)
+        >>> output[0]  # Will be zeros
+        Array([0., 0., 0.], dtype=float32)
+
+    Use max_norm to constrain embedding norms:
+
+    .. code-block:: python
+
+        >>> embedding = bst.nn.Embedding(num_embeddings=10, embedding_size=3, max_norm=1.0)
+        >>> # All embeddings accessed in a forward pass are renormalized to have norm <= 1.0
+
+    Load pretrained embeddings:
+
+    .. code-block:: python
+
+        >>> import brainstate
+        >>> import jax.numpy as jnp
+        >>> pretrained = jnp.array([[1.0, 2.0, 3.0],
+        ...                         [4.0, 5.0, 6.0],
+        ...                         [7.0, 8.0, 9.0]])
+        >>> embedding = bst.nn.Embedding.from_pretrained(pretrained, param_type=brainstate.FakeState)
+        >>> embedding.weight.value.shape
+        (3, 3)
+    """
+
+    def __init__(
+        self,
+        num_embeddings: int,
+        embedding_size: Size,
+        embedding_init: Union[Callable, ArrayLike] = init.LecunUniform(),
+        padding_idx: Optional[int] = None,
+        max_norm: Optional[float] = None,
+        norm_type: float = 2.0,
+        scale_grad_by_freq: bool = False,
+        freeze: bool = False,
+        name: Optional[str] = None,
+        param_type: type = ParamState,
+    ):
+        super().__init__(name=name)
+
+        self.num_embeddings = int(num_embeddings)
+        if self.num_embeddings < 0:
+            raise ValueError('num_embeddings must not be negative.')
+
+        embedding_size_tuple = _normalize_embedding_size(embedding_size)
+        self.embedding_size = embedding_size_tuple
+        self.out_size = embedding_size_tuple
+
+        if padding_idx is not None:
+            padding_idx = int(padding_idx)
+            if padding_idx < 0 or padding_idx >= self.num_embeddings:
+                raise ValueError(f'padding_idx must be within [0, {self.num_embeddings}).')
+        self.padding_idx = padding_idx
+
+        if max_norm is not None and max_norm <= 0:
+            raise ValueError('max_norm must be positive when provided.')
+        self.max_norm = max_norm
+        self.norm_type = norm_type
+        self.scale_grad_by_freq = bool(scale_grad_by_freq)
+        self.freeze = bool(freeze)
+
+        weight_shape = (self.num_embeddings, *self.out_size)
+        weight = init.param(embedding_init, weight_shape)
+
+        if self.padding_idx is not None:
+            weight = weight.at[self.padding_idx].set(0)
+
+        self.weight = param_type(weight)
+        self._lookup = _embedding_lookup_fn(self.padding_idx, self.scale_grad_by_freq)
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        embeddings: ArrayLike,
+        padding_idx: Optional[int] = None,
+        max_norm: Optional[float] = None,
+        norm_type: float = 2.0,
+        scale_grad_by_freq: bool = False,
+        freeze: bool = True,
+        name: Optional[str] = None,
+        param_type: type = ParamState,
+    ):
+        r"""
+        Create an Embedding instance from given 2-dimensional array.
+
+        Parameters
+        ----------
+        embeddings : ArrayLike
+            Array containing weights for the Embedding. First dimension is passed to
+            Embedding as ``num_embeddings``, remaining dimensions as ``embedding_size``.
+        padding_idx : int, optional
+            If specified, the entries at ``padding_idx`` do not contribute to the gradient.
+            Default is ``None``.
+        max_norm : float, optional
+            See module initialization documentation. Default is ``None``.
+        norm_type : float, optional
+            See module initialization documentation. Default is ``2.0``.
+        scale_grad_by_freq : bool, optional
+            See module initialization documentation. Default is ``False``.
+        freeze : bool, optional
+            If ``True``, embeddings are frozen (no gradients). Default is ``True``.
+        name : str, optional
+            The name of the module.
+
+        Returns
+        -------
+        Embedding
+            An Embedding module with pretrained weights.
+
+        Examples
+        --------
+        Load pretrained word embeddings:
+
+        .. code-block:: python
+
+            >>> import jax.numpy as jnp
+            >>> import brainstate as bst
+            >>> pretrained = jnp.array([[1.0, 2.0, 3.0],
+            ...                         [4.0, 5.0, 6.0],
+            ...                         [7.0, 8.0, 9.0]])
+            >>> embedding = bst.nn.Embedding.from_pretrained(pretrained)
+            >>> embedding.weight.value.shape
+            (3, 3)
+            >>> indices = jnp.array([1])
+            >>> embedding(indices)
+            Array([[4., 5., 6.]], dtype=float32)
+        """
+        embeddings = jnp.asarray(embeddings)
+        if embeddings.ndim < 2:
+            raise ValueError('embeddings must be at least 2-dimensional')
+
+        num_embeddings = embeddings.shape[0]
+        embedding_size = embeddings.shape[1:]
+
+        instance = cls(
+            num_embeddings=num_embeddings,
+            embedding_size=embedding_size,
+            embedding_init=embeddings,
+            padding_idx=padding_idx,
+            max_norm=max_norm,
+            norm_type=norm_type,
+            scale_grad_by_freq=scale_grad_by_freq,
+            freeze=freeze,
+            name=name,
+            param_type=param_type,
+        )
+
+        instance.weight = param_type(jnp.asarray(embeddings))
+        return instance
+
+    def update(self, indices: ArrayLike):
+        """
+        Retrieve embeddings for the given indices.
+
+        Parameters
+        ----------
+        indices : ArrayLike
+            Indices to retrieve embeddings for. Can be any shape.
+
+        Returns
+        -------
+        ArrayLike
+            Embeddings corresponding to the indices, with shape
+            ``(*indices.shape, *embedding_size)``.
+        """
+        indices = jnp.asarray(indices)
+        if not jnp.issubdtype(indices.dtype, jnp.integer):
+            raise TypeError('Embedding indices must be integers.')
+
+        weight_value = self.weight.value
+        effective_weight = weight_value
+
+        if self.max_norm is not None:
+            renormed_weight = self._apply_max_norm(weight_value, indices)
+            effective_weight = weight_value + jax.lax.stop_gradient(renormed_weight - weight_value)
+            if not _contains_tracer(renormed_weight):
+                self.weight.value = renormed_weight
+
+        if self.freeze:
+            effective_weight = jax.lax.stop_gradient(effective_weight)
+
+        embeddings = self._lookup(effective_weight, indices)
+        return embeddings
+
+    def _apply_max_norm(self, weight: jax.Array, indices: jax.Array) -> jax.Array:
+        """Apply max_norm constraint to the embedding weights for the given indices."""
+        flat_idx = jnp.ravel(indices)
+        if flat_idx.size == 0:
+            return weight
+
+        flat_idx = jnp.asarray(flat_idx, dtype=jnp.int32)
+        if self.padding_idx is not None:
+            pad_value = jnp.asarray(self.padding_idx, dtype=flat_idx.dtype)
+            flat_idx = flat_idx[flat_idx != pad_value]
+
+        if flat_idx.size == 0:
+            return weight
+
+        rows = weight[flat_idx]
+        rows_flat = rows.reshape((rows.shape[0], -1))
+        row_dtype = rows_flat.dtype
+
+        norms = jnp.linalg.norm(rows_flat, ord=self.norm_type, axis=1, keepdims=True)
+        max_norm = jnp.asarray(self.max_norm, dtype=row_dtype)
+        eps = jnp.asarray(1e-8, dtype=row_dtype)
+        one = jnp.asarray(1.0, dtype=row_dtype)
+        scale = jnp.minimum(one, max_norm / (norms + eps))
+        rows_scaled = (rows_flat * scale).reshape(rows.shape)
+
+        return weight.at[flat_idx].set(rows_scaled)
+
+

brainstate/nn/_embedding_test.py

@@ -0,0 +1,156 @@
+# Copyright 2024 BrainX Ecosystem Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+# -*- coding: utf-8 -*-
+
+import unittest
+
+import jax
+import jax.numpy as jnp
+
+import brainstate as bs
+
+
+class TestEmbedding(unittest.TestCase):
+    """Comprehensive tests for the Embedding module."""
+
+    def setUp(self):
+        settings = bs.environ.all()
+        self._prev_fit = settings.get('fit', None)
+        bs.environ.set(fit=True)
+
+    def tearDown(self):
+        if self._prev_fit is None:
+            bs.environ.pop('fit', None)
+        else:
+            bs.environ.set(fit=self._prev_fit)
+
+    def test_padding_idx_zero_gradient(self):
+        embedding = bs.nn.Embedding(num_embeddings=4, embedding_size=3, padding_idx=0)
+        lookup = embedding._lookup
+        weight = jnp.arange(12.0, dtype=jnp.float32).reshape(4, 3)
+        indices = jnp.array([0, 1, 0, 2], dtype=jnp.int32)
+
+        def loss_fn(w):
+            return jnp.sum(lookup(w, indices))
+
+        grad = jax.grad(loss_fn)(weight)
+        self.assertTrue(jnp.allclose(grad[0], 0.0))
+        self.assertFalse(jnp.allclose(grad[1], 0.0))
+
+    def test_scale_grad_by_freq(self):
+        base = bs.nn.Embedding(num_embeddings=5, embedding_size=2)
+        scaled = bs.nn.Embedding(num_embeddings=5, embedding_size=2, scale_grad_by_freq=True)
+        base_lookup = base._lookup
+        scaled_lookup = scaled._lookup
+
+        weight = jnp.arange(10.0, dtype=jnp.float32).reshape(5, 2)
+        indices = jnp.array([1, 1, 2], dtype=jnp.int32)
+
+        def loss_base(w):
+            return jnp.sum(base_lookup(w, indices))
+
+        def loss_scaled(w):
+            return jnp.sum(scaled_lookup(w, indices))
+
+        base_grad = jax.grad(loss_base)(weight)
+        scaled_grad = jax.grad(loss_scaled)(weight)
+
+        counts = jnp.bincount(indices, length=weight.shape[0])
+        ones = jnp.ones((indices.shape[0], weight.shape[1]), dtype=weight.dtype)
+        expected_base = jnp.zeros_like(weight).at[indices].add(ones)
+        expected_scaled = jnp.where(counts[:, None] > 0, jnp.ones_like(weight), 0.0)
+
+        self.assertTrue(jnp.allclose(base_grad, expected_base))
+        self.assertTrue(jnp.allclose(scaled_grad, expected_scaled))
+
+    def test_lookup_grad_jit_consistent(self):
+        embedding = bs.nn.Embedding(num_embeddings=4, embedding_size=2)
+        lookup = embedding._lookup
+        weight = jnp.arange(8.0, dtype=jnp.float32).reshape(4, 2)
+        indices = jnp.array([0, 1, 1, 3], dtype=jnp.int32)
+
+        def loss_fn(w):
+            return jnp.sum(lookup(w, indices))
+
+        grad_eager = jax.grad(loss_fn)(weight)
+        grad_jitted = jax.grad(jax.jit(loss_fn))(weight)
+
+        expected = jnp.zeros_like(weight).at[indices].add(jnp.ones((indices.shape[0], weight.shape[1]), dtype=weight.dtype))
+
+        self.assertTrue(jnp.allclose(grad_eager, grad_jitted))
+        self.assertTrue(jnp.allclose(grad_eager, expected))
+
+    def test_jit_forward_with_max_norm(self):
+        embedding = bs.nn.Embedding(num_embeddings=3, embedding_size=3, max_norm=0.5)
+        heavy = jnp.array([[0.0, 0.0, 0.0], [1.0, 2.0, 2.0], [0.3, -0.4, 0.5]], dtype=jnp.float32)
+        embedding.weight.value = heavy
+        indices = jnp.array([1, 2, 1], dtype=jnp.int32)
+
+        compiled = jax.jit(lambda ids: embedding(ids))
+        out = compiled(indices)
+        self.assertEqual(out.shape, (3, 3))
+        output_norms = jnp.linalg.norm(out, axis=-1)
+        self.assertTrue(jnp.all(output_norms <= 0.5 + 1e-6))
+        # Weight remains unclipped during JIT execution but must be usable without tracer leaks
+        self.assertGreater(float(jnp.linalg.norm(embedding.weight.value[1])), 0.5)
+
+    def test_freeze_disables_gradients(self):
+        embedding = bs.nn.Embedding(num_embeddings=4, embedding_size=2, freeze=True)
+        indices = jnp.array([1, 2, 3], dtype=jnp.int32)
+
+        def loss_fn(weight):
+            embedding.weight.value = weight
+            return jnp.sum(embedding(indices))
+
+        grad = jax.grad(loss_fn)(embedding.weight.value)
+        self.assertTrue(jnp.allclose(grad, 0.0))
+
+    def test_from_pretrained_defaults_to_freeze(self):
+        pretrained = jnp.arange(12.0, dtype=jnp.float32).reshape(4, 3)
+        embedding = bs.nn.Embedding.from_pretrained(pretrained)
+        self.assertTrue(embedding.freeze)
+
+        def loss_fn(weight):
+            embedding.weight.value = weight
+            return jnp.sum(embedding(jnp.array([1, 2], dtype=jnp.int32)))
+
+        grad = jax.grad(loss_fn)(embedding.weight.value)
+        self.assertTrue(jnp.allclose(grad, 0.0))
+
+    def test_from_pretrained_unfrozen_gradients(self):
+        pretrained = jnp.arange(6.0, dtype=jnp.float32).reshape(2, 3)
+        embedding = bs.nn.Embedding.from_pretrained(pretrained, freeze=False)
+
+        def loss_fn(weight):
+            embedding.weight.value = weight
+            return jnp.sum(embedding(jnp.array([0, 1], dtype=jnp.int32)))
+
+        grad = jax.grad(loss_fn)(embedding.weight.value)
+        self.assertFalse(jnp.allclose(grad, 0.0))
+
+    def test_max_norm_renormalizes_weights(self):
+        embedding = bs.nn.Embedding(num_embeddings=3, embedding_size=3, max_norm=1.0, norm_type=2.0)
+        custom_weight = jnp.array([[0.0, 0.0, 0.0], [3.0, 4.0, 0.0], [0.5, 0.5, 0.5]], dtype=jnp.float32)
+        embedding.weight.value = custom_weight
+        _ = embedding(jnp.array([1, 2], dtype=jnp.int32))
+
+        row_norm = jnp.linalg.norm(embedding.weight.value[1])
+        self.assertLessEqual(float(row_norm), 1.0 + 1e-6)
+        self.assertTrue(jnp.allclose(embedding.weight.value[0], custom_weight[0]))
+
+
+if __name__ == '__main__':
+    unittest.main()