Equimo 0.1.0__py3-none-any.whl

equimo/layers/ffn.py

@@ -0,0 +1,184 @@
+from typing import Callable
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+import jax.random as jr
+from jaxtyping import Array, Float, PRNGKeyArray
+
+from equimo.layers.dropout import Dropout
+
+
+class Mlp(eqx.Module):
+    """Multi-layer Perceptron (MLP) module with dropout.
+
+    A standard MLP implementation with two fully connected layers, activation function,
+    and dropout for regularization. The architecture follows:
+    input -> fc1 -> activation -> dropout1 -> fc2 -> dropout2 -> output
+
+    Attributes:
+        fc1: First linear layer
+        fc2: Second linear layer
+        drop1: Dropout after first layer
+        drop2: Dropout after second layer
+        act_layer: Activation function
+    """
+
+    fc1: eqx.nn.Linear
+    fc2: eqx.nn.Linear
+    drop1: Dropout
+    drop2: Dropout
+    act_layer: Callable
+
+    def __init__(
+        self,
+        in_features: int,
+        *,
+        key: PRNGKeyArray,
+        out_features: int | None = None,
+        hidden_features: int | None = None,
+        act_layer: Callable = jax.nn.gelu,
+        dropout_rate: float = 0.0,
+        bias: bool = True,
+        **kwargs,
+    ):
+        """Initialize the MLP.
+
+        Args:
+            in_features: Number of input features
+            key: PRNG key for initialization
+            out_features: Number of output features (default: same as in_features)
+            hidden_features: Number of hidden features (default: same as in_features)
+            act_layer: Activation function (default: gelu)
+            dropout_rate: Dropout probability (default: 0.0)
+            bias: Whether to include bias in linear layers (default: True)
+            **kwargs: Additional arguments
+        """
+        key_fc1, key_fc2 = jr.split(key, 2)
+
+        self.act_layer = act_layer
+
+        hidden_features = hidden_features or in_features
+        out_features = out_features or in_features
+
+        self.fc1 = eqx.nn.Linear(
+            in_features, hidden_features, use_bias=bias, key=key_fc1
+        )
+        self.fc2 = eqx.nn.Linear(
+            hidden_features, out_features, use_bias=bias, key=key_fc2
+        )
+
+        self.drop1 = Dropout(dropout_rate)
+        self.drop2 = Dropout(dropout_rate)
+
+    def __call__(
+        self,
+        x: Float[Array, "seqlen dim"],
+        enable_dropout: bool,
+        key: PRNGKeyArray,
+    ) -> Float[Array, "seqlen dim"]:
+        key_dr1, key_dr2 = jr.split(key, 2)
+
+        x = self.drop1(
+            self.act_layer(jax.vmap(self.fc1)(x)),
+            inference=not enable_dropout,
+            key=key_dr1,
+        )
+        x = self.drop2(
+            jax.vmap(self.fc2)(x),
+            inference=not enable_dropout,
+            key=key_dr2,
+        )
+
+        return x
+
+
+class SwiGlu(eqx.Module):
+    """SwiGLU activation module with dropout.
+
+    Implements the SwiGLU (Swish-Gated Linear Unit) activation function with dropout,
+    as described in "GLU Variants Improve Transformer" paper [1]. The architecture uses
+    a gating mechanism where the input is transformed by two parallel paths and
+    combined multiplicatively.
+
+    The computation flow is:
+    1. Joint projection to higher dimension (w12)
+    2. Split into two paths
+    3. Apply SiLU to first path and multiply with second path
+    4. Project back to original dimension (w3)
+
+    Attributes:
+        w12: Joint projection layer for both paths
+        w3: Final projection layer
+        drop1: Dropout after gating
+        drop2: Dropout after final projection
+
+    References:
+        [1]: https://arxiv.org/pdf/2002.05202
+    """
+
+    w12: eqx.nn.Linear
+    w3: eqx.nn.Linear
+    drop1: Dropout
+    drop2: Dropout
+
+    def __init__(
+        self,
+        in_features: int,
+        *,
+        key: PRNGKeyArray,
+        out_features: int | None = None,
+        hidden_features: int | None = None,
+        dropout_rate: float = 0.0,
+        bias: bool = True,
+        **kwargs,
+    ):
+        """Initialize the SwiGLU module.
+
+        Args:
+            in_features: Number of input features
+            key: PRNG key for initialization
+            out_features: Number of output features (default: same as in_features)
+            hidden_features: Size of hidden dimension (default: same as in_features)
+            dropout_rate: Dropout probability (default: 0.0)
+            bias: Whether to include bias in linear layers (default: True)
+            **kwargs: Additional arguments
+        """
+        key_fc1, key_fc2 = jr.split(key, 2)
+
+        hidden_features = hidden_features or in_features
+        out_features = out_features or in_features
+
+        self.w12 = eqx.nn.Linear(
+            in_features, hidden_features, use_bias=bias, key=key_fc1
+        )
+        self.w3 = eqx.nn.Linear(
+            hidden_features // 2, out_features, use_bias=bias, key=key_fc2
+        )
+
+        self.drop1 = Dropout(dropout_rate)
+        self.drop2 = Dropout(dropout_rate)
+
+    def __call__(
+        self,
+        x: Float[Array, "seqlen dim"],
+        enable_dropout: bool,
+        key: PRNGKeyArray,
+    ) -> Float[Array, "seqlen dim"]:
+        key_dr1, key_dr2 = jr.split(key, 2)
+
+        x12 = jax.vmap(self.w12)(x)
+        x1, x2 = jnp.split(x12, 2, axis=-1)
+        x = self.drop1(
+            jax.nn.silu(x1) * x2,
+            inference=not enable_dropout,
+            key=key_dr1,
+        )
+
+        x = self.drop2(
+            jax.vmap(self.w3)(x),
+            inference=not enable_dropout,
+            key=key_dr2,
+        )
+
+        return x

equimo/layers/generic.py

@@ -0,0 +1,68 @@
+import equinox as eqx
+from jaxtyping import Array, Float, PRNGKeyArray
+
+from equimo.layers.dropout import DropPathAdd
+
+
+class Residual(eqx.Module):
+    """A wrapper module that adds a residual connection with optional drop path.
+
+    This module wraps any other module and adds a residual (skip) connection around it.
+    It also includes drop path regularization which stochastically drops the residual
+    path during training. The computation flow is:
+    input -> [main branch: module] + [residual branch: identity with drop path] -> output
+
+    Attributes:
+        module: The module to wrap with a residual connection
+        drop_path: DropPath module for residual connection regularization
+    """
+
+    module: eqx.Module
+    drop_path: DropPathAdd
+
+    def __init__(
+        self,
+        module: eqx.Module,
+        drop_path: float = 0,
+    ):
+        """Initialize the Residual wrapper.
+
+        Args:
+            module: The module to wrap with a residual connection
+            drop_path: Drop path rate (probability of dropping the residual connection)
+                      (default: 0)
+        """
+        self.module = module
+        self.drop_path = DropPathAdd(drop_path)
+
+    def __call__(
+        self,
+        x: Float[Array, "..."],
+        enable_dropout: bool,
+        key: PRNGKeyArray,
+        pass_args: bool = False,
+    ) -> Float[Array, "..."]:
+        """Forward pass of the residual block.
+
+        Args:
+            x: Input tensor of any shape
+            enable_dropout: Whether to enable dropout during training
+            key: PRNG key for randomness
+            pass_args: Whether to pass enable_dropout and key to the wrapped module
+                      (default: False)
+
+        Returns:
+            Output tensor with same shape as input, combining the module output
+            with the residual connection through drop path
+        """
+        if pass_args:
+            x2 = self.module(x, enable_dropout=enable_dropout, key=key)
+        else:
+            x2 = self.module(x)
+
+        return self.drop_path(
+            x,
+            x2,
+            inference=not enable_dropout,
+            key=key,
+        )

equimo/layers/mamba.py

@@ -0,0 +1,169 @@
+import math
+from typing import List, Tuple
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+import jax.random as jr
+from einops import rearrange
+from jaxtyping import Array, Float, PRNGKeyArray
+
+from equimo.layers.norm import RMSNormGated
+from equimo.ops.scan import non_causal_linear_attn
+
+
+class Mamba2Mixer(eqx.Module):
+    """Mamba2 Mixer.
+
+    This class implements the a Mamba2 Mixer using State Space Duality (SSD),
+    from Mamba2 [1]. Also supports implementation details from Visual State Space
+    Duality (VSSD) [2].
+
+    Attributes:
+        in_proj (eqx.nn.Linear): Input projection layer.
+        conv (eqx.nn.Conv): Convolutional layer for processing input.
+        dt_bias (eqx.nn.Param): Bias for delta time.
+        A_log (eqx.nn.Param): Logarithm of the state transition matrix A.
+        D (eqx.nn.Param): Direct feedthrough matrix D.
+        norm (eqx.nn.RMSNorm): Root Mean Square Layer Normalization.
+        out_proj (eqx.nn.Linear): Output projection layer.
+
+    Args:
+        config (MambaConfig): Configuration object for the Mamba2VisionMixer.
+        rngs (nnx.Rngs): Random number generators for parameter initialization.
+
+    Notes:
+        This implementation is heavily based on wlln/scratch.
+
+    References:
+        [1] Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality
+            (https://arxiv.org/abs/2401.04054)
+        [2] VSSD: Vision Mamba with Non-Causal State Space Duality
+            (https://arxiv.org/abs/2407.18559)
+    """
+
+    d_inner: int = eqx.field(static=True)
+    n_heads: int = eqx.field(static=True)
+    head_dim: int = eqx.field(static=True)
+    n_groups: int = eqx.field(static=True)
+    indices_xBC: List[int] = eqx.field(static=True)
+
+    in_proj: eqx.nn.Linear
+    conv: eqx.nn.Conv
+    dt_bias: Float[Array, "n_heads"]
+    A_log: Float[Array, "n_heads"]
+    D: Float[Array, "n_heads"]
+    norm: eqx.nn.LayerNorm
+    out_proj: eqx.nn.Linear
+
+    def __init__(
+        self,
+        d_model: int,
+        *,
+        key: PRNGKeyArray,
+        expand: int = 2,
+        n_groups: int = 1,
+        head_dim: int = 64,
+        d_state: int = 128,
+        d_conv: int = 4,
+        dt_min: float = 0.001,
+        dt_max: float = 0.1,
+        dt_init_floor: float = 1e-4,
+        A_init_range: Tuple[int, int] = (1, 16),
+        use_bias: bool = False,
+        conv_bias: bool = True,
+        **kwargs,
+    ):
+        key_inproj, key_outproj, key_conv, key_randvals, key_a = jr.split(key, 5)
+        self.d_inner = int(d_model * expand)
+        if self.d_inner % head_dim != 0:
+            raise ValueError("`d_inner` must be a multiple of `head_dim`.")
+        self.n_heads = self.d_inner // head_dim
+        self.head_dim = head_dim
+        self.n_groups = n_groups
+        self.indices_xBC = [self.d_inner, self.d_inner + n_groups * d_state]
+
+        d_in_proj = 2 * self.d_inner + 2 * n_groups * d_state + self.n_heads
+        self.in_proj = eqx.nn.Linear(
+            d_model,
+            d_in_proj,
+            use_bias=use_bias,
+            key=key_inproj,
+        )
+
+        conv_dim = self.d_inner + 2 * n_groups * d_state
+        self.conv = eqx.nn.Conv(
+            num_spatial_dims=1,
+            in_channels=conv_dim,
+            out_channels=conv_dim,
+            kernel_size=d_conv,
+            groups=conv_dim,
+            padding="SAME",
+            use_bias=conv_bias,
+            key=key_conv,
+        )
+
+        rand_vals = jr.uniform(key_randvals, (self.n_heads,))
+        dt = jnp.exp(
+            rand_vals * (math.log(dt_max) - math.log(dt_min)) + math.log(dt_min)
+        )
+        dt = jnp.clip(dt, a_min=dt_init_floor)
+        self.dt_bias = dt + jnp.log(-jnp.expm1(-dt))
+
+        A_min, A_max = A_init_range
+        A = jr.uniform(key_a, (self.n_heads,), minval=A_min, maxval=A_max)
+        self.A_log = jnp.log(A)
+
+        self.D = jnp.ones(self.n_heads)
+
+        self.norm = eqx.nn.LayerNorm(self.d_inner)
+        self.out_proj = eqx.nn.Linear(
+            self.d_inner, d_model, use_bias=use_bias, key=key_outproj
+        )
+
+    def __call__(
+        self,
+        x: Float[Array, "seqlen dim"],
+        enable_dropout: bool,
+        key: PRNGKeyArray,
+    ) -> Float[Array, "seqlen dim"]:
+        A = -jnp.exp(self.A_log)
+        zxbcdt = jax.vmap(self.in_proj)(x)
+
+        z, xbc, dt = jnp.split(
+            zxbcdt,
+            [self.d_inner, zxbcdt.shape[-1] - self.n_heads],
+            axis=-1,
+        )
+
+        dt = jax.nn.softplus(dt + self.dt_bias)
+
+        # Pad or truncate the xbc tensor to match the conv kernel size
+        # xbc_rearranged = rearrange(xbc, "b l d -> b d l")
+        # conv_state = pad_or_truncate_to_length(xbc_rearranged, self.config.d_conv)
+
+        # apply 1d convolution and silu activation
+        xbc_conv = rearrange(self.conv(rearrange(xbc, "s d -> d s")), "d s -> s d")
+        xbc_silu = jax.nn.silu(xbc_conv[: x.shape[0], :])
+
+        # split the conv state into the conv kernel and the conv state
+        x, B, C = jnp.split(xbc_silu, self.indices_xBC, axis=-1)
+
+        x = rearrange(x, "l (h p) -> l h p", p=self.head_dim)
+
+        y = non_causal_linear_attn(
+            x, dt=dt, A=A, B=B, C=C, D=self.D, n_groups=self.n_groups
+        )
+
+        y = rearrange(y, "l h p -> l (h p)")
+
+        # apply the output projection
+        if isinstance(self.norm, RMSNormGated):
+            y = self.norm(y, jax.nn.silu(z))
+        else:
+            # Should be LayerNorm
+            y = jax.vmap(self.norm)(y) * z
+
+        y = jax.vmap(self.out_proj)(y)
+
+        return y

equimo/layers/norm.py

@@ -0,0 +1,87 @@
+from typing import Optional
+
+import equinox as eqx
+import jax.numpy as jnp
+from jax import lax
+from jaxtyping import Array, Float
+
+
+class RMSNormGated(eqx.Module):
+    """Root Mean Square (RMS) Normalization with optional gating.
+
+    Implements RMS normalization with learnable scale parameters and optional
+    gating mechanism. RMS norm is similar to Layer Norm but only normalizes by
+    the root mean square, without centering the mean.
+
+    Attributes:
+        w: Learnable scale parameter vector of size dim
+    """
+
+    w: Float[Array, "dim"]
+
+    def __init__(self, d: int):
+        """Initialize RMSNormGated.
+
+        Args:
+            d: Dimension of the input features
+        """
+        self.w = jnp.ones(d)
+
+    def __call__(
+        self,
+        x: Float[Array, "dim"],
+        z: Optional[Float[Array, "dim"]] = None,
+    ) -> Float[Array, "dim"]:
+        """Apply RMS normalization with optional gating.
+
+        Args:
+            x: Input tensor of shape (dim,)
+            z: Optional gating tensor of shape (dim,)
+
+        Returns:
+            Normalized tensor of same shape as input
+        """
+        if z is not None:
+            x *= z
+
+        y = x.astype(jnp.float32)
+        norm = y * lax.rsqrt(jnp.mean(y * y, -1, keepdims=True) + 1e-5)
+
+        return self.w * norm.astype(x.dtype)
+
+
+class LayerScale(eqx.Module):
+    """Layer scaling module for stabilizing deep networks.
+
+    Implements learnable scaling factors for each feature dimension, initialized
+    to a small value. This helps stabilize training in deep networks by initially
+    dampening the contribution of each layer.
+
+    Attributes:
+        init_values: Initial scale value (static)
+        gamma: Learnable scale parameters of size dim
+    """
+
+    init_values: float = eqx.field(static=True)
+    gamma: Optional[Float[Array, "dim"]]
+
+    def __init__(self, dim: int, init_values: float):
+        """Initialize LayerScale.
+
+        Args:
+            dim: Dimension of the input features
+            init_values: Initial value for all scaling factors
+        """
+        self.init_values = init_values
+        self.gamma = jnp.repeat(self.init_values, dim)
+
+    def __call__(self, x: Float[Array, "dim"]):
+        """Apply layer scaling to input tensor.
+
+        Args:
+            x: Input tensor of shape (dim,)
+
+        Returns:
+            Scaled tensor of same shape as input
+        """
+        return x * self.gamma