matnets 1.1.0__py3-none-any.whl

matnets/__init__.py

@@ -0,0 +1,9 @@
+"""JAX primitives for matrix-valued neural networks."""
+
+from matnets import lax, nn
+from matnets._dense import dense
+from matnets._params import MatrixParams, init
+
+__all__ = ["MatrixParams", "dense", "init", "lax", "nn"]
+
+__version__ = "0.1.0"

matnets/_dense.py

@@ -0,0 +1,76 @@
+"""Core dense matrix primitive."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+import jax
+import jax.numpy as jnp
+from jax import Array
+
+from matnets._params import MatrixParams
+
+
+def identity(x: Array) -> Array:
+    return x
+
+
+@jax.custom_jvp
+def _dense_linear(W: Array, B: Array, x: Array) -> Array:
+    q, p, a, k = W.shape
+    c = x.shape[-1]
+    W_flat = jnp.reshape(jnp.transpose(W, (0, 2, 1, 3)), (q * a, p * k))
+    x_flat = jnp.reshape(x, (p * k, c))
+    out = jnp.matmul(W_flat, x_flat)
+    return jnp.reshape(out, (q, a, c)) + B
+
+
+@_dense_linear.defjvp
+def _dense_linear_jvp(
+    primals: tuple[Array, Array, Array],
+    tangents: tuple[Array, Array, Array],
+) -> tuple[Array, Array]:
+    W, B, x = primals
+    dW, dB, dx = tangents
+    q, p, a, k = W.shape
+    c = x.shape[-1]
+
+    W_flat = jnp.reshape(jnp.transpose(W, (0, 2, 1, 3)), (q * a, p * k))
+    dW_flat = jnp.reshape(jnp.transpose(dW, (0, 2, 1, 3)), (q * a, p * k))
+    x_flat = jnp.reshape(x, (p * k, c))
+    dx_flat = jnp.reshape(dx, (p * k, c))
+
+    y_flat = jnp.matmul(W_flat, x_flat)
+    y = jnp.reshape(y_flat, (q, a, c)) + B
+
+    dy_flat = jnp.matmul(dW_flat, x_flat) + jnp.matmul(W_flat, dx_flat)
+    dy = jnp.reshape(dy_flat, (q, a, c)) + dB
+
+    return y, dy
+
+
+def dense(
+    params: MatrixParams,
+    x: Array,
+    activation: Callable[[Array], Array] = identity,
+) -> Array:
+    """Apply the core square-matrix contraction ``qpak,pkc -> qac``."""
+
+    W = params.W
+    x = jnp.asarray(x)
+    if W.ndim != 4:
+        msg = "dense expects weights shaped (q, p, n, n)"
+        raise ValueError(msg)
+
+    q, p, n, k = W.shape
+    if n != k:
+        msg = "dense weights must map square matrices: last two axes must match"
+        raise ValueError(msg)
+    if x.shape != (p, n, n):
+        msg = f"dense expects input shaped ({p}, {n}, {n})"
+        raise ValueError(msg)
+    if params.B.shape != (q, n, n):
+        msg = f"dense bias must be shaped ({q}, {n}, {n})"
+        raise ValueError(msg)
+
+    return activation(_dense_linear(W, params.B, x))

matnets/_params.py

@@ -0,0 +1,66 @@
+"""Parameter containers and initializers for matrix primitives."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from math import sqrt
+
+import jax
+import jax.numpy as jnp
+from jax import Array
+from jax.typing import DTypeLike
+
+
+@jax.tree_util.register_pytree_node_class
+@dataclass(frozen=True)
+class MatrixParams:
+    """Weights and bias for matrix contractions.
+
+    Dense weights use ``(q, p, n, n)`` and consume inputs shaped ``(p, n, n)``.
+    The bias is shaped ``(q, n, n)`` so each output unit has a full matrix bias.
+    Structural primitives may use higher-rank weights while keeping the same
+    small, pytree-friendly container.
+    """
+
+    W: Array
+    B: Array
+
+    def tree_flatten(self) -> tuple[tuple[Array, Array], None]:
+        return (self.W, self.B), None
+
+    @classmethod
+    def tree_unflatten(
+        cls,
+        aux_data: None,
+        children: tuple[Array, Array],
+    ) -> MatrixParams:
+        del aux_data
+        return cls(*children)
+
+
+def init(
+    key: Array,
+    p: int,
+    q: int,
+    n: int,
+    *,
+    dtype: DTypeLike = jnp.float32,
+) -> MatrixParams:
+    """Initialize dense matrix parameters with Glorot-uniform weights."""
+
+    if p <= 0 or q <= 0 or n <= 0:
+        msg = "p, q, and n must all be positive"
+        raise ValueError(msg)
+
+    fan_in = p * n
+    fan_out = q * n
+    limit = sqrt(6.0 / (fan_in + fan_out))
+    W = jax.random.uniform(
+        key,
+        shape=(q, p, n, n),
+        minval=-limit,
+        maxval=limit,
+        dtype=dtype,
+    )
+    B = jnp.zeros((q, n, n), dtype=dtype)
+    return MatrixParams(W=W, B=B)

matnets/lax/__init__.py

@@ -0,0 +1,16 @@
+"""JAX-level structural primitives."""
+
+from matnets.lax.attention import (
+    frobenius_score,
+    matrix_attention,
+    scaled_frobenius_score,
+)
+from matnets.lax.conv import matrix_conv1d, matrix_conv2d
+
+__all__ = [
+    "frobenius_score",
+    "matrix_attention",
+    "matrix_conv1d",
+    "matrix_conv2d",
+    "scaled_frobenius_score",
+]

matnets/lax/attention.py

@@ -0,0 +1,62 @@
+"""Attention primitive for matrix-valued tokens."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from math import sqrt
+
+import jax
+import jax.numpy as jnp
+from jax import Array
+
+from matnets._dense import dense
+from matnets._params import MatrixParams
+
+ScoreFn = Callable[[Array, Array], Array]
+
+
+def frobenius_score(Q: Array, K: Array) -> Array:
+    """Scalar Frobenius inner product between two matrix-valued tokens."""
+
+    return jnp.einsum("pkc,pkc->", Q, K)
+
+
+def scaled_frobenius_score(Q: Array, K: Array) -> Array:
+    scale = sqrt(Q.shape[0] * Q.shape[1] * Q.shape[2])
+    return frobenius_score(Q, K) / scale
+
+
+def matrix_attention(
+    params: MatrixParams | None,
+    Q: Array,
+    K: Array,
+    V: Array,
+    score_fn: ScoreFn = scaled_frobenius_score,
+) -> Array:
+    """Score, normalize, aggregate, and optionally project matrix tokens.
+
+    ``Q``, ``K``, and ``V`` are token sequences shaped ``(t, p, n, n)``.
+    ``score_fn`` defines the pairwise semantics and must return a scalar.
+    Passing ``params=None`` returns the aggregated context directly; otherwise
+    each context token is projected through ``dense(params, token)``.
+    """
+
+    Q = jnp.asarray(Q)
+    K = jnp.asarray(K)
+    V = jnp.asarray(V)
+    if Q.ndim != 4 or K.ndim != 4 or V.ndim != 4:
+        msg = "matrix_attention expects Q, K, and V shaped (t, p, n, n)"
+        raise ValueError(msg)
+    if Q.shape[1:] != K.shape[1:] or Q.shape[1:] != V.shape[1:]:
+        msg = "Q, K, and V must share neuron and matrix axes"
+        raise ValueError(msg)
+    if Q.shape[-2] != Q.shape[-1]:
+        msg = "matrix_attention expects square matrix tokens"
+        raise ValueError(msg)
+
+    scores = jax.vmap(lambda q: jax.vmap(lambda k: score_fn(q, k))(K))(Q)
+    weights = jax.nn.softmax(scores, axis=-1)
+    context = jnp.einsum("ij,jpkc->ipkc", weights, V)
+    if params is None:
+        return context
+    return jax.vmap(lambda token: dense(params, token))(context)

matnets/lax/conv.py

@@ -0,0 +1,153 @@
+"""Convolution-like structural matrix primitives."""
+
+from __future__ import annotations
+
+from typing import Literal
+
+import jax
+import jax.numpy as jnp
+from jax import Array
+
+from matnets._params import MatrixParams
+
+Padding = Literal["VALID", "SAME"]
+
+
+def _check_bias_shape(params: MatrixParams, q: int, n: int) -> None:
+    if params.B.shape != (q, n, n):
+        msg = f"bias must be shaped ({q}, {n}, {n})"
+        raise ValueError(msg)
+
+
+def _check_square_matrix_kernel(W: Array, *, rank: int, shape_name: str) -> None:
+    if W.ndim != rank:
+        msg = f"{shape_name} expects weights shaped with rank {rank}"
+        raise ValueError(msg)
+    if W.shape[-2] != W.shape[-1]:
+        msg = (
+            "matrix convolution weights must map square matrices: "
+            "last two axes must match"
+        )
+        raise ValueError(msg)
+
+
+def _positive_stride(stride: int) -> int:
+    if stride <= 0:
+        msg = "stride must be positive"
+        raise ValueError(msg)
+    return stride
+
+
+def _same_padding(size: int) -> tuple[int, int]:
+    left = (size - 1) // 2
+    right = size - 1 - left
+    return left, right
+
+
+def matrix_conv1d(
+    params: MatrixParams,
+    x: Array,
+    *,
+    stride: int = 1,
+    padding: Padding = "VALID",
+) -> Array:
+    """Slide matrix kernels across a sequence.
+
+    ``params.W`` is shaped ``(q, p, r, n, n)`` and ``x`` is shaped
+    ``(t, p, n, n)``. The output is ``(t_out, q, n, n)``.
+    """
+
+    stride = _positive_stride(stride)
+    W = params.W
+    _check_square_matrix_kernel(W, rank=5, shape_name="matrix_conv1d")
+
+    q, p, kernel, n, _ = W.shape
+    x = jnp.asarray(x)
+    if x.ndim != 4:
+        msg = "matrix_conv1d expects inputs shaped (t, p, n, n)"
+        raise ValueError(msg)
+    if x.shape[1:] != (p, n, n):
+        msg = "input axes must match weight axes: x must be shaped (t, p, n, n)"
+        raise ValueError(msg)
+    _check_bias_shape(params, q, n)
+
+    if padding == "SAME":
+        pad_left, pad_right = _same_padding(kernel)
+    elif padding == "VALID":
+        pad_left, pad_right = 0, 0
+    else:
+        msg = "padding must be 'VALID' or 'SAME'"
+        raise ValueError(msg)
+
+    x_reshaped = jnp.reshape(x, (x.shape[0], p * n, n))
+
+    W_trans = jnp.transpose(W, (0, 3, 1, 4, 2))
+    W_reshaped = jnp.reshape(W_trans, (q * n, p * n, kernel))
+
+    out_conv = jax.lax.conv_general_dilated(
+        x_reshaped,
+        W_reshaped,
+        window_strides=(stride,),
+        padding=((pad_left, pad_right),),
+        dimension_numbers=("WCN", "OIW", "WCN"),
+    )
+    return jnp.reshape(out_conv, (out_conv.shape[0], q, n, n)) + params.B
+
+
+def matrix_conv2d(
+    params: MatrixParams,
+    x: Array,
+    *,
+    stride: int | tuple[int, int] = 1,
+    padding: Padding = "VALID",
+) -> Array:
+    """Slide matrix kernels across a 2D grid.
+
+    ``params.W`` is shaped ``(q, p, h, w, n, n)`` and ``x`` is shaped
+    ``(y, x, p, n, n)``. The output is ``(y_out, x_out, q, n, n)``.
+    """
+
+    if isinstance(stride, int):
+        stride_y = stride_x = _positive_stride(stride)
+    else:
+        stride_y = _positive_stride(stride[0])
+        stride_x = _positive_stride(stride[1])
+
+    W = params.W
+    _check_square_matrix_kernel(W, rank=6, shape_name="matrix_conv2d")
+
+    q, p, kernel_y, kernel_x, n, _ = W.shape
+    x = jnp.asarray(x)
+    if x.ndim != 5:
+        msg = "matrix_conv2d expects inputs shaped (y, x, p, n, n)"
+        raise ValueError(msg)
+    if x.shape[2:] != (p, n, n):
+        msg = "input axes must match weight axes: x must be shaped (y, x, p, n, n)"
+        raise ValueError(msg)
+    _check_bias_shape(params, q, n)
+
+    if padding == "SAME":
+        pad_top, pad_bottom = _same_padding(kernel_y)
+        pad_left, pad_right = _same_padding(kernel_x)
+    elif padding == "VALID":
+        pad_top = pad_bottom = pad_left = pad_right = 0
+    else:
+        msg = "padding must be 'VALID' or 'SAME'"
+        raise ValueError(msg)
+
+    x_reshaped = jnp.reshape(x, (x.shape[0], x.shape[1], p * n, n))
+
+    W_trans = jnp.transpose(W, (0, 4, 1, 5, 2, 3))
+    W_reshaped = jnp.reshape(W_trans, (q * n, p * n, kernel_y, kernel_x))
+
+    out_conv = jax.lax.conv_general_dilated(
+        x_reshaped,
+        W_reshaped,
+        window_strides=(stride_y, stride_x),
+        padding=((pad_top, pad_bottom), (pad_left, pad_right)),
+        dimension_numbers=("HWCN", "OIHW", "HWCN"),
+    )
+    return (
+        jnp.reshape(out_conv, (out_conv.shape[0], out_conv.shape[1], q, n, n))
+        + params.B
+    )

matnets/nn/__init__.py

@@ -0,0 +1,5 @@
+"""Neural-network patterns composed from matrix primitives."""
+
+from matnets.nn.recurrent import gru_step, lstm_step, rnn_step
+
+__all__ = ["gru_step", "lstm_step", "rnn_step"]

matnets/nn/recurrent.py

@@ -0,0 +1,60 @@
+"""Recurrent patterns built from the dense matrix primitive."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Mapping
+
+import jax
+import jax.numpy as jnp
+from jax import Array
+
+from matnets._dense import dense
+from matnets._params import MatrixParams
+
+
+def rnn_step(
+    params: MatrixParams,
+    carry: Array,
+    x: Array,
+    *,
+    activation: Callable[[Array], Array] = jnp.tanh,
+) -> tuple[Array, Array]:
+    """Simple RNN step suitable for use inside ``jax.lax.scan``."""
+
+    combined = jnp.concatenate([carry, x], axis=0)
+    next_carry = dense(params, combined, activation)
+    return next_carry, next_carry
+
+
+def lstm_step(
+    params: Mapping[str, MatrixParams],
+    carry: tuple[Array, Array],
+    x: Array,
+) -> tuple[tuple[Array, Array], Array]:
+    """LSTM step using ``i``, ``f``, ``g``, and ``o`` dense gate params."""
+
+    h, c = carry
+    combined = jnp.concatenate([h, x], axis=0)
+    i = dense(params["i"], combined, jax.nn.sigmoid)
+    f = dense(params["f"], combined, jax.nn.sigmoid)
+    g = dense(params["g"], combined, jnp.tanh)
+    o = dense(params["o"], combined, jax.nn.sigmoid)
+    next_c = f * c + i * g
+    next_h = o * jnp.tanh(next_c)
+    return (next_h, next_c), next_h
+
+
+def gru_step(
+    params: Mapping[str, MatrixParams],
+    carry: Array,
+    x: Array,
+) -> tuple[Array, Array]:
+    """GRU step using ``z``, ``r``, and ``n`` dense gate params."""
+
+    combined = jnp.concatenate([carry, x], axis=0)
+    z = dense(params["z"], combined, jax.nn.sigmoid)
+    r = dense(params["r"], combined, jax.nn.sigmoid)
+    candidate_input = jnp.concatenate([r * carry, x], axis=0)
+    n = dense(params["n"], candidate_input, jnp.tanh)
+    next_carry = (1.0 - z) * n + z * carry
+    return next_carry, next_carry

matnets/py.typed

@@ -0,0 +1 @@
+

matnets-1.1.0.dist-info/METADATA

@@ -0,0 +1,54 @@
+Metadata-Version: 2.4
+Name: matnets
+Version: 1.1.0
+Summary: A small experimental neural network library where neurons are represented as matrices.
+Project-URL: Homepage, https://github.com/dsainvg/MATNETS
+Project-URL: Documentation, https://github.com/dsainvg/MATNETS/tree/main/docs
+Project-URL: Issues, https://github.com/dsainvg/MATNETS/issues
+Author: dsainvg
+License-Expression: MIT
+License-File: LICENSE
+Keywords: jax,machine-learning,matrix,neural-networks
+Classifier: Development Status :: 2 - Pre-Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.11
+Requires-Dist: jax>=0.4
+Provides-Extra: dev
+Requires-Dist: mypy>=1.10; extra == 'dev'
+Requires-Dist: pytest>=8.0; extra == 'dev'
+Requires-Dist: ruff>=0.5; extra == 'dev'
+Provides-Extra: docs
+Requires-Dist: mkdocs-material>=9.5; extra == 'docs'
+Requires-Dist: mkdocs>=1.6; extra == 'docs'
+Description-Content-Type: text/markdown
+
+# MATNETS
+
+MATNETS is a small JAX library for matrix-neuron neural network experiments.
+Each neuron carries an `n x n` matrix instead of a scalar.
+
+The user documentation lives in [`docs/`](docs/index.md):
+
+- [`docs/index.md`](docs/index.md): overview
+- [`docs/getting-started.md`](docs/getting-started.md): install and first model
+- [`docs/concepts.md`](docs/concepts.md): matrix-neuron shapes and JAX transforms
+- [`docs/api.md`](docs/api.md): API guide
+- [`docs/examples.md`](docs/examples.md): runnable examples
+- [`docs/development.md`](docs/development.md): local development commands
+
+Quick check:
+
+```powershell
+.\.venv\Scripts\python.exe examples\five_hidden_net.py
+.\.venv\Scripts\python.exe -m pytest
+```
+
+MIT license.

matnets-1.1.0.dist-info/RECORD

@@ -0,0 +1,13 @@
+matnets/__init__.py,sha256=s90bc4_V_TVyLgnmf1JbD-MiqWEGEZ9rSAUMMsHVdQM,246
+matnets/_dense.py,sha256=00Riy9Nkqm4U7qo4wAS-xCiNeBnotNrvvs0z6kNRYTQ,2061
+matnets/_params.py,sha256=wKdDyW-St-VxMyhSzgAGuDG5muhYTCTYa_nLrF3CY0s,1610
+matnets/py.typed,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+matnets/lax/__init__.py,sha256=7zw_HdUXRqUmfFY9pYITNQEWQRSy25DZq00FjEFD2pw,341
+matnets/lax/attention.py,sha256=ImdhAusgzdv_6BUNscePcqWHPEV2IXgb7BczYrFWd10,1979
+matnets/lax/conv.py,sha256=wP5Ly-hHpFes9jXepPcJrzG_xli7sMnb3Ak6CTT4vz4,4484
+matnets/nn/__init__.py,sha256=KFp0s_hxBsOm59vGp4yGYg0zccaA0GwRi3KA6qcrBg0,176
+matnets/nn/recurrent.py,sha256=Vtk0BSnc-OEUWx1uUyjmIdn2GRqtAihCSDxZaFQ9grs,1767
+matnets-1.1.0.dist-info/METADATA,sha256=f90gfRdkQQR7KG74rcGLXcXk5psn2XLfD8QnLGkp3Cw,2061
+matnets-1.1.0.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
+matnets-1.1.0.dist-info/licenses/LICENSE,sha256=Y2Z4Zzs_BrNYRx8nWVtax4k14klf8SELDd77VK0wui8,1077
+matnets-1.1.0.dist-info/RECORD,,

matnets-1.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.29.0
+Root-Is-Purelib: true
+Tag: py3-none-any

matnets-1.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 MATNETS contributors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.