matnets 1.1.0__tar.gz

matnets-1.1.0/.github/workflows/docs.yml

@@ -0,0 +1,63 @@
+name: Docs
+
+on:
+  push:
+    branches:
+      - main
+    paths:
+      - "docs/**"
+      - "mkdocs.yml"
+      - "pyproject.toml"
+      - ".github/workflows/docs.yml"
+  workflow_dispatch:
+
+permissions:
+  contents: read
+  pages: write
+  id-token: write
+
+concurrency:
+  group: pages
+  cancel-in-progress: false
+
+jobs:
+  build:
+    name: Build MkDocs site
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Check out repository
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+          cache: pip
+
+      - name: Configure GitHub Pages
+        uses: actions/configure-pages@v5
+
+      - name: Install docs dependencies
+        run: python -m pip install -e ".[docs]"
+
+      - name: Build site
+        run: mkdocs build --strict
+
+      - name: Upload Pages artifact
+        uses: actions/upload-pages-artifact@v3
+        with:
+          path: site
+
+  deploy:
+    name: Deploy to GitHub Pages
+    needs: build
+    runs-on: ubuntu-latest
+    environment:
+      name: github-pages
+      url: "${{ steps.deployment.outputs.page_url }}"
+
+    steps:
+      - name: Deploy site
+        id: deployment
+        uses: actions/deploy-pages@v4

matnets-1.1.0/.github/workflows/pypi.yaml

@@ -0,0 +1,30 @@
+# pypi.yaml
+# Configuration for PyPI publishing (example for GitHub Actions)
+# Adjust as needed for your CI/CD provider and secrets management
+
+name: Publish Python 🐍 distribution 📦 to PyPI
+
+on:
+  push:
+    tags:
+      - 'v*.*.*'
+
+jobs:
+  build-and-publish:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.11'
+      - name: Install build dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install build
+      - name: Build package
+        run: python -m build
+      - name: Publish package to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          password: ${{ secrets.PYPI_API_TOKEN }}

matnets-1.1.0/.github/workflows/testing.yml

@@ -0,0 +1,82 @@
+name: CI
+
+on:
+  pull_request:
+    paths-ignore:
+      - "docs/**"
+      - "mkdocs.yml"
+      - ".github/workflows/docs.yml"
+      - "*.md"
+  push:
+    branches:
+      - main
+    paths-ignore:
+      - "docs/**"
+      - "mkdocs.yml"
+      - ".github/workflows/docs.yml"
+      - "*.md"
+
+jobs:
+  test:
+    name: Test Python ${{ matrix.python-version }}
+    runs-on: ubuntu-latest
+
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version:
+          - "3.11"
+          - "3.12"
+          - "3.13"
+          - "3.14"
+
+    steps:
+      - name: Check out repository
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+          cache: pip
+
+      - name: Install package
+        run: python -m pip install -e ".[dev,docs]"
+
+      - name: Lint
+        run: ruff check src/ tests/
+
+      - name: Type check
+        run: mypy src
+
+      - name: Test
+        run: pytest
+
+      - name: Build docs
+        run: mkdocs build --strict
+
+  package:
+    name: Build package
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Check out repository
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.13"
+          cache: pip
+
+      - name: Install build tools
+        run: python -m pip install build
+
+      - name: Build distribution
+        run: python -m build
+
+      - name: Upload distribution artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: python-package
+          path: dist/*

matnets-1.1.0/.gitignore

@@ -0,0 +1,15 @@
+.vscode/
+__pycache__/
+*.py[cod]
+*.egg-info/
+.pytest_cache/
+.ruff_cache/
+.mypy_cache/
+.venv/
+venv/
+build/
+dist/
+site/
+.coverage
+htmlcov/
+runs/

matnets-1.1.0/CHANGELOG.md

@@ -0,0 +1,6 @@
+# Changelog
+
+## 0.1.0
+
+- Initial project scaffold.
+- Added matrix-neuron, layer, and sequential network primitives.

matnets-1.1.0/CONTRIBUTING.md

@@ -0,0 +1,20 @@
+# Contributing
+
+Thanks for helping MATNETS grow.
+
+## Local Setup
+
+```bash
+python -m pip install -e ".[dev,docs]"
+```
+
+## Checks
+
+```bash
+pytest
+ruff check .
+mypy src
+```
+
+Keep changes focused, add tests for behavior changes, and update docs when the
+public API changes.

matnets-1.1.0/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.

matnets-1.1.0/PKG-INFO

@@ -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/README.md

@@ -0,0 +1,22 @@
+# 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/docs/api.md

@@ -0,0 +1,164 @@
+# API Guide
+
+## `matnets.MatrixParams`
+
+`MatrixParams` stores the weights and bias for matrix primitives:
+
+```python
+from matnets import MatrixParams
+```
+
+Dense parameters use:
+
+```text
+W: (q, p, n, n)
+B: (q, n, n)
+```
+
+`MatrixParams` is registered as a JAX pytree, so it works with `jax.jit`,
+`jax.vmap`, `jax.grad`, and nested dictionaries/lists of parameters.
+
+## `matnets.init`
+
+```python
+params = matnets.init(key, p=2, q=3, n=4)
+```
+
+Creates:
+
+```text
+params.W: (3, 2, 4, 4)
+params.B: (3, 4, 4)
+```
+
+Weights use Glorot-uniform initialization. Bias starts at zero.
+
+## `matnets.dense`
+
+```python
+y = matnets.dense(params, x)
+```
+
+Expected shapes:
+
+```text
+params.W: (q, p, n, n)
+params.B: (q, n, n)
+x:        (p, n, n)
+y:        (q, n, n)
+```
+
+With activation:
+
+```python
+y = matnets.dense(params, x, activation=jax.nn.relu)
+```
+
+The core operation is:
+
+```python
+jnp.einsum("qpak,pkc->qac", params.W, x) + params.B
+```
+
+## `matnets.lax.matrix_conv1d`
+
+```python
+from matnets.lax import matrix_conv1d
+
+y = matrix_conv1d(params, x, stride=1, padding="VALID")
+```
+
+Expected shapes:
+
+```text
+params.W: (q, p, r, n, n)
+params.B: (q, n, n)
+x:        (t, p, n, n)
+y:        (t_out, q, n, n)
+```
+
+`r` is the 1D kernel size.
+
+## `matnets.lax.matrix_conv2d`
+
+```python
+from matnets.lax import matrix_conv2d
+
+y = matrix_conv2d(params, x, stride=(1, 1), padding="SAME")
+```
+
+Expected shapes:
+
+```text
+params.W: (q, p, h, w, n, n)
+params.B: (q, n, n)
+x:        (height, width, p, n, n)
+y:        (height_out, width_out, q, n, n)
+```
+
+## `matnets.lax.matrix_attention`
+
+```python
+from matnets.lax import matrix_attention
+
+out = matrix_attention(None, Q, K, V)
+```
+
+Expected token shapes:
+
+```text
+Q:   (tokens_q, p, n, n)
+K:   (tokens_k, p, n, n)
+V:   (tokens_k, p, n, n)
+out: (tokens_q, p, n, n)
+```
+
+By default the score is a scaled Frobenius inner product. You can pass a custom
+`score_fn` that receives one query token and one key token and returns a scalar.
+
+If `params` is not `None`, each aggregated output token is projected through
+`matnets.dense(params, token)`.
+
+## `matnets.nn`
+
+`matnets.nn` contains recurrent wiring patterns built from `dense`.
+
+```python
+from matnets.nn import rnn_step, lstm_step, gru_step
+```
+
+These functions are intended to be used with `jax.lax.scan`.
+
+### RNN
+
+```python
+carry, outputs = jax.lax.scan(
+    lambda h, x_t: rnn_step(params, h, x_t),
+    h0,
+    sequence,
+)
+```
+
+### LSTM
+
+```python
+carry, outputs = jax.lax.scan(
+    lambda carry, x_t: lstm_step(params, carry, x_t),
+    (h0, c0),
+    sequence,
+)
+```
+
+LSTM params must contain keys `"i"`, `"f"`, `"g"`, and `"o"`.
+
+### GRU
+
+```python
+carry, outputs = jax.lax.scan(
+    lambda h, x_t: gru_step(params, h, x_t),
+    h0,
+    sequence,
+)
+```
+
+GRU params must contain keys `"z"`, `"r"`, and `"n"`.

matnets-1.1.0/docs/concepts.md

@@ -0,0 +1,78 @@
+# Concepts
+
+## Matrix-Neurons
+
+A traditional dense layer usually maps vectors:
+
+```text
+x: (p)
+W: (q, p)
+y: (q)
+```
+
+MATNETS maps stacks of square matrices:
+
+```text
+x: (p, n, n)
+W: (q, p, n, n)
+B: (q, n, n)
+y: (q, n, n)
+```
+
+`p` is the input neuron count. `q` is the output neuron count. `n` is the
+matrix size inside each neuron.
+
+## Dense Primitive
+
+The core operation is:
+
+```python
+jnp.einsum("qpak,pkc->qac", W, x) + B
+```
+
+Under the square-matrix contract:
+
+```text
+a == n
+k == n
+c == n
+```
+
+so the output is always `(q, n, n)`.
+
+## Bias
+
+The bias is a full matrix:
+
+```text
+B: (q, n, n)
+```
+
+Each output matrix-neuron gets its own complete matrix bias.
+
+## JAX Transforms
+
+MATNETS functions are ordinary JAX functions. You can transform them with:
+
+```python
+jax.jit(forward)
+jax.vmap(forward, in_axes=(None, 0))
+jax.grad(loss)
+jax.lax.scan(step, carry, sequence)
+```
+
+The main parallel work is the dense einsum. `vmap` adds batch or token axes
+around it. `scan` handles recurrence over time while each step still uses
+compiled dense contractions.
+
+## Recurrent State
+
+RNN, LSTM, and GRU hidden states are stacks of matrices:
+
+```text
+H: (hidden_neurons, n, n)
+C: (hidden_neurons, n, n)
+```
+
+Gates are also matrix-valued, so an LSTM forget gate has one value per matrix
+entry, not just one scalar per neuron.

matnets-1.1.0/docs/development.md

@@ -0,0 +1,42 @@
+# Development
+
+## Test
+
+```powershell
+.\.venv\Scripts\python.exe -m pytest
+```
+
+## Run Examples
+
+```powershell
+.\.venv\Scripts\python.exe examples\basic_forward.py
+.\.venv\Scripts\python.exe examples\five_hidden_net.py
+.\.venv\Scripts\python.exe examples\matrix_architectures.py
+```
+
+## Syntax Check
+
+```powershell
+.\.venv\Scripts\python.exe -m compileall -q src examples tests
+```
+
+## Build Documentation
+
+```powershell
+.\.venv\Scripts\python.exe -m pip install -e ".[docs]"
+.\.venv\Scripts\python.exe -m mkdocs build --strict
+```
+
+## Package Layout
+
+```text
+src/matnets/
+  _params.py          MatrixParams and init()
+  _dense.py           core dense matrix primitive
+  lax/conv.py         convolution-like matrix primitives
+  lax/attention.py    matrix attention primitive
+  nn/recurrent.py     RNN, LSTM, and GRU step patterns
+tests/                test suite
+examples/             runnable examples
+docs/                 user documentation
+```

matnets-1.1.0/docs/examples.md

@@ -0,0 +1,52 @@
+# Examples
+
+## Basic Forward Pass
+
+```powershell
+.\.venv\Scripts\python.exe examples\basic_forward.py
+```
+
+This creates one dense layer and applies it to an input shaped `(2, 2, 2)`.
+
+## Five Hidden Layers
+
+```powershell
+.\.venv\Scripts\python.exe examples\five_hidden_net.py
+```
+
+This example defines a small class:
+
+```python
+model = FiveHiddenNet(key, input_neurons=3, hidden_neurons=4, n=2)
+y = jax.jit(model.forward)(model.params, x)
+```
+
+The output shape is `(1, 2, 2)`.
+
+## Architecture Walkthrough
+
+```powershell
+.\.venv\Scripts\python.exe examples\matrix_architectures.py
+```
+
+This file checks shape flow through:
+
+- MLP
+- batched MLP with `jax.vmap`
+- gradients with `jax.grad`
+- RNN with `jax.lax.scan`
+- LSTM with `jax.lax.scan`
+- Frobenius attention
+- residual block
+
+## Where Parallelization Happens
+
+The dense operation:
+
+```python
+jnp.einsum("qpak,pkc->qac", W, x)
+```
+
+is the main kernel. JAX can compile it with `jit`, map it over batches or token
+sequences with `vmap`, differentiate it with `grad`, and call it repeatedly
+inside `lax.scan`.