jaxonlayers 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl

jaxonlayers/functions/embedding.py

@@ -0,0 +1,21 @@
+import jax.numpy as jnp
+from jaxtyping import Array, Float, Int
+
+
+def sinusoidal_embedding(
+    t: Int[Array, ""], embedding_size: int
+) -> Float[Array, " embedding_size"]:
+    if embedding_size % 2 != 0:
+        raise ValueError(f"Embedding size must be even, but got {embedding_size}")
+
+    half_dim = embedding_size // 2
+    embedding_freqs = jnp.exp(
+        -jnp.log(10000)
+        * jnp.arange(start=0, stop=half_dim, dtype=jnp.float32)
+        / half_dim
+    )
+
+    time_args = t * embedding_freqs
+    embedding = jnp.concatenate([jnp.sin(time_args), jnp.cos(time_args)])
+
+    return embedding

jaxonlayers/functions/state_space.py

@@ -4,27 +4,76 @@ from jaxtyping import Array, Float
 
 
 def selective_scan(
-    x: Float[Array, "seq_length d_inner"],
+    u: Float[Array, "seq_length d_inner"],
     delta: Float[Array, "seq_length d_inner"],
     A: Float[Array, "d_inner d_state"],
-    B: Float[Array, "seq_length d_state"],
-    C: Float[Array, "seq_length d_state"],
+    B: Float[Array, "seq_length d_inner d_state"],
+    C: Float[Array, "seq_length d_inner d_state"],
     D: Float[Array, " d_inner"],
+    chunk_size: int = 128,
 ) -> Float[Array, "seq_length d_inner"]:
-    L, d_inner = x.shape
-    _, d_state = A.shape
-    delta_A = jnp.exp(jnp.einsum("l d,d n -> l d n", delta, A))
-    delta_B_u = jnp.einsum("l d,l n,l d -> l d n", delta, B, x)
+    deltaA = jnp.exp(jnp.einsum("l d, d n -> l d n", delta, A))
+    deltaB_u = jnp.einsum("l d, l d n, l d -> l d n", delta, B, u)
 
-    x_res = jnp.zeros(shape=(d_inner, d_state))
+    seq_len, d_inner = u.shape
+    d_state = A.shape[1]
 
-    def step(x, i):
-        x = delta_A[i] * x + delta_B_u[i]
+    num_chunks = (seq_len + chunk_size - 1) // chunk_size
+    padded_len = num_chunks * chunk_size
 
-        y = jnp.einsum("d n,n -> d", x, C[i, :])
-        return x, y
+    pad_len = padded_len - seq_len
+    deltaA_padded = jnp.pad(deltaA, ((0, pad_len), (0, 0), (0, 0)))
+    deltaB_u_padded = jnp.pad(deltaB_u, ((0, pad_len), (0, 0), (0, 0)))
+    C_padded = jnp.pad(C, ((0, pad_len), (0, 0), (0, 0)))
 
-    _, ys = jax.lax.scan(step, x_res, jnp.arange(L))
+    deltaA_chunked = deltaA_padded.reshape(num_chunks, chunk_size, d_inner, d_state)
+    deltaB_u_chunked = deltaB_u_padded.reshape(num_chunks, chunk_size, d_inner, d_state)
+    C_chunked = C_padded.reshape(num_chunks, chunk_size, d_inner, d_state)
 
-    ys = ys + x * D
-    return ys
+    def intra_chunk_step(h_prev, scan_inputs):
+        deltaA_i, deltaB_u_i, C_i = scan_inputs
+        h_i = deltaA_i * h_prev + deltaB_u_i
+        y_i = jnp.einsum("d n, d n -> d", h_i, C_i)
+        return h_i, y_i
+
+    h0 = jnp.zeros((d_inner, d_state))
+
+    _, y_chunks = jax.vmap(jax.lax.scan, in_axes=(None, None, 0))(
+        intra_chunk_step, h0, (deltaA_chunked, deltaB_u_chunked, C_chunked)
+    )
+
+    def inter_chunk_step(carry_prev, scan_inputs):
+        A_prev, h_prev = carry_prev
+        deltaA_i, deltaB_u_i = scan_inputs
+
+        A_new = deltaA_i * A_prev
+        h_new = deltaA_i * h_prev + deltaB_u_i
+
+        return (A_new, h_new), (A_new, h_new)
+
+    A_carry_initial = jnp.ones((d_inner, d_state))
+    h_carry_initial = jnp.zeros((d_inner, d_state))
+    initial_carry = (A_carry_initial, h_carry_initial)
+
+    scan_inputs = (deltaA_chunked[:, -1], deltaB_u_chunked[:, -1])
+
+    _, (A_carry, h_carry) = jax.lax.scan(inter_chunk_step, initial_carry, scan_inputs)
+
+    A_carry = jnp.roll(A_carry, 1, axis=0)
+    h_carry = jnp.roll(h_carry, 1, axis=0)
+    A_carry = A_carry.at[0].set(jnp.ones_like(A_carry[0]))
+    h_carry = h_carry.at[0].set(jnp.zeros_like(h_carry[0]))
+
+    h_carry_broadcast = jnp.expand_dims(h_carry, axis=1)
+    h_correction = deltaA_chunked * h_carry_broadcast
+    y_carry = jnp.einsum("csdn, csdn -> csd", C_chunked, h_correction)
+
+    y_final = y_chunks + y_carry
+
+    y_final = y_final.reshape(padded_len, d_inner)
+
+    y_unpadded = y_final[:seq_len]
+
+    output = y_unpadded.real + u * D
+
+    return output

jaxonlayers/functions/utils.py

@@ -1,5 +1,7 @@
+import equinox as eqx
 import jax
 import jax.numpy as jnp
+from jaxtyping import PyTree
 
 
 def default_floating_dtype():
@@ -7,3 +9,50 @@ def default_floating_dtype():
         return jnp.float64
     else:
         return jnp.float32
+
+
+def summarize_model(model: PyTree) -> str:
+    params, _ = eqx.partition(model, eqx.is_array)
+
+    param_counts = {}
+    total_params = 0
+
+    def count_params(pytree, name=""):
+        nonlocal total_params
+        count = 0
+        if isinstance(pytree, jnp.ndarray):
+            count = pytree.size
+            total_params += count
+            if name:
+                param_counts[name] = count
+        elif hasattr(pytree, "__dict__"):
+            for key, value in pytree.__dict__.items():
+                subname = f"{name}.{key}" if name else key
+                count += count_params(value, subname)
+        elif isinstance(pytree, (list, tuple)):
+            for i, value in enumerate(pytree):
+                subname = f"{name}[{i}]" if name else f"[{i}]"
+                count += count_params(value, subname)
+        elif isinstance(pytree, dict):
+            for key, value in pytree.items():
+                subname = f"{name}.{key}" if name else str(key)
+                count += count_params(value, subname)
+        return count
+
+    count_params(params)
+
+    # Display as table
+    lines = []
+    lines.append("Model Parameter Summary")
+    lines.append("=" * 50)
+    lines.append(f"{'Parameter Name':<30} {'Count':<15}")
+    lines.append("-" * 50)
+
+    for name, count in param_counts.items():
+        lines.append(f"{name:<30} {count:<15,}")
+
+    lines.append("-" * 50)
+    lines.append(f"{'Total Parameters':<30} {total_params:<15,}")
+    lines.append("=" * 50)
+
+    return "\n".join(lines)

jaxonlayers/layers/attention.py

@@ -38,8 +38,8 @@ class MultiheadAttention(eqx.Module):
 
     def __init__(
         self,
-        embed_dim,
-        num_heads,
+        embed_dim: int,
+        num_heads: int,
         dropout=0.0,
         bias=True,
         add_bias_kv=False,

jaxonlayers/layers/state_space.py

@@ -11,14 +11,17 @@ class SelectiveStateSpace(eqx.Module):
     input_proj: eqx.nn.Linear
     delta_proj: eqx.nn.Linear
     A_log: Float[Array, "d_inner d_state"]
-    D: Float[Array, " d_inner"]
+    D: Float[Array, "d_inner"]
+    out_proj: eqx.nn.Linear
 
     d_inner: int = eqx.field(static=True)
     dt_rank: int = eqx.field(static=True)
     d_state: int = eqx.field(static=True)
+    d_model: int = eqx.field(static=True)
 
     def __init__(
         self,
+        d_model: int,
         d_inner: int,
         dt_rank: int,
         d_state: int,
@@ -31,19 +34,19 @@ class SelectiveStateSpace(eqx.Module):
         if dtype is None:
             dtype = default_floating_dtype()
         assert dtype is not None
+
+        self.d_model = d_model
         self.d_inner = d_inner
         self.dt_rank = dt_rank
         self.d_state = d_state
-        (
-            key,
-            input_proj_key,
-            delta_proj_key,
-        ) = jax.random.split(key, 3)
+
+        keys = jax.random.split(key, 4)
+        proj_dim = self.dt_rank + 2 * self.d_inner * self.d_state
         self.input_proj = eqx.nn.Linear(
-            d_inner,
-            dt_rank + d_state * 2,
+            self.d_model,
+            proj_dim,
             use_bias=use_input_proj_bias,
-            key=input_proj_key,
+            key=keys[0],
             dtype=dtype,
         )
 
@@ -51,25 +54,37 @@ class SelectiveStateSpace(eqx.Module):
             dt_rank,
             d_inner,
             use_bias=use_delta_proj_bias,
-            key=delta_proj_key,
+            key=keys[1],
             dtype=dtype,
         )
-        A = jnp.repeat(jnp.arange(1, d_state + 1), d_inner).reshape(d_inner, d_state)
+
+        A = jnp.arange(1, d_state + 1, dtype=jnp.float32)
+        A = jnp.tile(A, (d_inner, 1))
         self.A_log = jnp.log(A)
+
         self.D = jnp.ones(d_inner, dtype=dtype)
 
+        self.out_proj = eqx.nn.Linear(
+            d_inner, d_model, use_bias=False, key=keys[2], dtype=dtype
+        )
+
     def __call__(self, x: Float[Array, "seq_length d_inner"]):
-        A = -jnp.exp(self.A_log)
-        D = self.D
+        L, _ = x.shape
+        A = -jnp.exp(self.A_log.astype(jnp.float32))
+        D = self.D.astype(jnp.float32)
 
         delta_b_c = jax.vmap(self.input_proj)(x)
+        delta, B, C = jnp.split(
+            delta_b_c,
+            [self.dt_rank, self.dt_rank + self.d_inner * self.d_state],
+            axis=-1,
+        )
+
+        B = B.reshape(L, self.d_inner, self.d_state)
+        C = C.reshape(L, self.d_inner, self.d_state)
 
-        split_indices = [
-            self.dt_rank,
-            self.dt_rank + self.d_state,
-        ]
-        delta, B, C = jnp.split(delta_b_c, split_indices, axis=-1)
         delta = jax.nn.softplus(jax.vmap(self.delta_proj)(delta))
 
         y = selective_scan(x, delta, A, B, C, D)
-        return y
+
+        return jax.vmap(self.out_proj)(y)

jaxonlayers-0.1.2.dist-info/METADATA

@@ -0,0 +1,16 @@
+Metadata-Version: 2.4
+Name: jaxonlayers
+Version: 0.1.2
+Summary: Add your description here
+Requires-Python: >=3.13
+Requires-Dist: beartype>=0.21.0
+Requires-Dist: equinox>=0.13.0
+Requires-Dist: jax>=0.8.0
+Requires-Dist: jaxtyping>=0.3.2
+Description-Content-Type: text/markdown
+
+# jaxonlayers
+
+This library provides some utility function and useful layers that extend the [Equinox](https://github.com/patrick-kidger/equinox) library.
+
+The aim was to create them to be the PyTorch equivalent, i.e. to match their PyTorch counterpart's output.

{jaxonlayers-0.1.0.dist-info → jaxonlayers-0.1.2.dist-info}/RECORD

@@ -1,21 +1,21 @@
 jaxonlayers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 jaxonlayers/functions/__init__.py,sha256=lxMNSnEn2pJ1XLLyQTTFXiCcJTdrzfE1RlorqD9lEog,711
 jaxonlayers/functions/attention.py,sha256=AyaAeA2yo5Cgljk2rU6JlVttQeG0FMBLE-f7285PdM0,12945
+jaxonlayers/functions/embedding.py,sha256=g76Wg_MSyqatwjWrhDSfAxDaMyr2s3A5AoST_cRa3Q8,601
 jaxonlayers/functions/initialization.py,sha256=h7uzdPl-rL7faT9hbRs5aN6EeaX0r70y88ad1cwZvmY,1161
 jaxonlayers/functions/masking.py,sha256=5VeHMuoVeKxAHtzGb74GqvKHfib4wJAETTAn1oaOgLM,1531
 jaxonlayers/functions/normalization.py,sha256=e2vNNbDz-Y6j5bgQshd4MshnzLcDOjUnGSinpRXtPtA,206
 jaxonlayers/functions/regularization.py,sha256=ltFVeJZFhdNrHsH0CjDl9Y3dlAQmh1ABntCjpYUBgXM,1317
-jaxonlayers/functions/state_space.py,sha256=_aIk4hXHZRtU2117z9qlfYjo-iAwA_dCkn9OpUEBtEk,831
-jaxonlayers/functions/utils.py,sha256=ahNtu8wf5A3wVjRUg6z5kPktDbEHqsak4sezRu6ap0E,184
+jaxonlayers/functions/state_space.py,sha256=_PdykvDSRZWmIQR5KZTUN6shLhPGElio9THnoVpLJ_g,2747
+jaxonlayers/functions/utils.py,sha256=M1uc01yOADy2ig6YsIyn0JBf_zdFuJS-rxin5RVub-A,1749
 jaxonlayers/layers/__init__.py,sha256=gjH0QUOCplii5KNYshWNehPfHZLVjEeYk8EhhY7FHRE,480
 jaxonlayers/layers/abstract.py,sha256=uyPKGsH5DETnjiU3PUiF98tk9boNZW96amxrM5JQlZY,366
-jaxonlayers/layers/attention.py,sha256=lAOIhzutQv1TcdECBJF8G95VoGxmPMBMe52S24H0nIo,7774
+jaxonlayers/layers/attention.py,sha256=NdjDnA3yYFpvXD2_SjBwF-8Pdk_GeRLPuPekTR-iMwE,7784
 jaxonlayers/layers/convolution.py,sha256=k0dMFBDjzycB7UNuyHqKihJtBa6u93V6OLxyUUyipN4,3247
 jaxonlayers/layers/normalization.py,sha256=3aGzNzDN05A72ZHLUM2w9WpicLtGsjzj1l0jhuyn63U,8379
 jaxonlayers/layers/regularization.py,sha256=ZrvtBJPH84xuxrxEbZc7TBxjp8OvKEv4ecan5s8F9zs,563
 jaxonlayers/layers/sequential.py,sha256=Tw98hNZiXMC-CYZD6h_pi7eAxkgHeQAUvZF2I9H0d8Y,2833
-jaxonlayers/layers/state_space.py,sha256=Nesj2Ts3mCCqE-u7PeB8roJbQXUql7rG0AIpqUVMqvg,2131
-jaxonlayers-0.1.0.dist-info/METADATA,sha256=02mfySFIYtuSbEYqaZUZqTY7op6jWHdICubD77Tr_Cg,275
-jaxonlayers-0.1.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-jaxonlayers-0.1.0.dist-info/top_level.txt,sha256=n5UHFDErh3dJY77ypkEKlwFOQffKBnpGH9nUrwCinto,12
-jaxonlayers-0.1.0.dist-info/RECORD,,
+jaxonlayers/layers/state_space.py,sha256=oDVRbduNtU48Q4rLd-XywZcqVN0QYTlq1UUhOXcGLoo,2537
+jaxonlayers-0.1.2.dist-info/METADATA,sha256=NGl7HoPwh1sC8JmgOeBD3kVEJgy5tmTR7jhlKM9um7k,539
+jaxonlayers-0.1.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+jaxonlayers-0.1.2.dist-info/RECORD,,

{jaxonlayers-0.1.0.dist-info → jaxonlayers-0.1.2.dist-info}/WHEEL

@@ -1,5 +1,4 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.9.0)
+Generator: hatchling 1.27.0
 Root-Is-Purelib: true
 Tag: py3-none-any
-

jaxonlayers-0.1.0.dist-info/METADATA

@@ -1,10 +0,0 @@
-Metadata-Version: 2.4
-Name: jaxonlayers
-Version: 0.1.0
-Summary: Add your description here
-Requires-Python: >=3.13
-Description-Content-Type: text/markdown
-Requires-Dist: beartype>=0.21.0
-Requires-Dist: equinox>=0.13.0
-Requires-Dist: jax>=0.7.2
-Requires-Dist: jaxtyping>=0.3.2

jaxonlayers-0.1.0.dist-info/top_level.txt

@@ -1 +0,0 @@
-jaxonlayers