Metadata-Version: 2.4
Name: jax-blox
Version: 0.1.6
Summary: A functional and lightweight neural network library for JAX.
Project-URL: Repository, https://github.com/hamzamerzic/blox
Author-email: Hamza Merzić <hamzamerzic@gmail.com>
License: MIT License
        
        Copyright (c) 2025 Hamza Merzić
        
        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.
License-File: LICENSE
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Education
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: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.11
Requires-Dist: chex>=0.1.0
Requires-Dist: jax>=0.10.0
Requires-Dist: jaxlib>=0.10.0
Requires-Dist: treescope
Provides-Extra: dev
Requires-Dist: pyink; extra == 'dev'
Requires-Dist: pyrefly; extra == 'dev'
Requires-Dist: pytest; extra == 'dev'
Requires-Dist: pytest-cov; extra == 'dev'
Requires-Dist: ruff; extra == 'dev'
Requires-Dist: tensorflow; extra == 'dev'
Provides-Extra: docs
Requires-Dist: furo; extra == 'docs'
Requires-Dist: ipykernel; extra == 'docs'
Requires-Dist: nbsphinx; extra == 'docs'
Requires-Dist: sphinx<9,>=7; extra == 'docs'
Description-Content-Type: text/markdown

<div align="center">
  <img src="https://i.ibb.co/FLmR2T3r/logo.png" width="400" alt="blox logo">

  <h1>blox</h1>

  <p>
    <strong>A functional and lightweight neural network library for JAX.</strong>
  </p>

  <a href="LICENSE"><img src="https://img.shields.io/badge/license-MIT-blue.svg" alt="blox is released under the MIT license"></a>
  <img src="https://img.shields.io/badge/python-3.11+-blue" alt="Python 3.11+">
  <img src="https://img.shields.io/badge/jax-0.8+-green" alt="JAX 0.8+">
</div>

---

**blox** unlocks the full potential of JAX by embracing its functional nature instead of fighting it.

JAX gives you composable transformations that let you write math and have it run blazingly fast on any hardware. **blox** is a thin layer on top that keeps all of that power accessible while giving you just enough structure to organize your neural networks.

Most JAX libraries try to force Object-Oriented paradigms to make JAX feel like PyTorch. While comfortable at first, this fights against JAX's functional nature. It introduces implicit global state and hidden contexts that eventually increase cognitive load and steepen the learning curve.

**blox** takes the opposite approach. We embrace the functional nature of JAX.

The entire mental model fits in one line:

```python
outputs, params = model(params, inputs)
```

Parameters go in, outputs and updated parameters come out. This is the standard pattern for [stateful computations in JAX](https://docs.jax.dev/en/latest/stateful-computations.html). Because state flows explicitly through your code, all JAX transformations—`jax.jit`, `jax.grad`, `jax.vmap`, `jax.checkpoint`—work out of the box. No wrappers, no decorators, no surprises.

## 🎯 Who is blox for?

* **Learners:** There is no "framework magic" to learn here. What you see is what you get. It is the best way to understand how neural networks actually work at the JAX level.
* **Practitioners:** If you're tired of fighting frameworks that hide important details, **blox** gives you complete transparency. Whether you're building custom training loops, implementing novel architectures, or scaling up, you have direct access to the full execution stack.

## 📦 Installation

Since blox uses JAX, check the [JAX installation guide](https://jax.readthedocs.io/en/latest/installation.html) for your specific hardware.

You will need Python 3.11 or later.

```bash
pip install jax-blox
```

## 🚀 The Basics

### Your First Layer

Let's build a linear layer to see how it feels. Notice the signature: `params` carries all model state, while `inputs` is your data.

```python
import jax
import jax.numpy as jnp
import blox as bx


class Linear(bx.Module):

  def __init__(self, graph: bx.Graph, output_size: int, rng: bx.Rng):
    super().__init__(graph)
    self.output_size = output_size
    self.rng = rng

  def __call__(self, params: bx.Params, x: jax.Array):
    # Parameters are created lazily on first use.
    # No need to specify input shapes upfront or preallocate memory!
    kernel, params = self.get_param(
        params,
        name='kernel',
        shape=(x.shape[-1], self.output_size),
        init=jax.nn.initializers.normal(),
        rng=self.rng,
    )
    bias, params = self.get_param(
        params,
        name='bias',
        shape=(self.output_size,),
        init=jax.nn.initializers.zeros,
        rng=self.rng,
    )
    return x @ kernel + bias, params
```

### Composing Layers

Modules are just Python objects. You can nest them, inject them, or generate them dynamically.

```python
class MLP(bx.Module):

  def __init__(
      self,
      graph: bx.Graph,
      hidden_size: int,
      output_size: int,
      rng: bx.Rng,
  ):
    super().__init__(graph)
    # graph.child('name') creates a unique path for each parameter.
    self.hidden = Linear(graph.child('hidden'), hidden_size, rng=rng)
    self.output = Linear(graph.child('output'), output_size, rng=rng)

  def __call__(self, params: bx.Params, x: jax.Array):
    x, params = self.hidden(params, x)
    x = jax.nn.relu(x)
    return self.output(params, x)
```

### Initialization & Inspection

We cleanly separate "Initialization" (traversing the graph to create parameters) from "Runtime".

```python
# Define the structure.
graph = bx.Graph('net')
rng = bx.Rng(graph.child('rng'))
model = MLP(graph.child('mlp'), hidden_size=128, output_size=10, rng=rng)

# Initialize the parameter container and initialize the RNG state (seed).
# We initialize the RNG first since we use it to initialize other modules.
params = bx.Params()
params = rng.seed(params, seed=42)

# Run a forward pass to trigger lazy parameter initialization.
dummy_input = jnp.ones((1, 784))
_, params = model(params, dummy_input)

# Lock it down to prevent accidental parameter creation during training.
params = params.locked()

# Visualize the full graph and parameter structure.
bx.display(graph, params)
```

## ⚡ JIT Compilation

**blox** modules are trivially compatible with `jax.jit`.

```python
# Just wrap and call. No special decorators needed.
outputs, params = jax.jit(model)(params, inputs)
```

## 📦 The Params Container

The `Params` container holds **all** model state in one place: weights, RNG state, batch norm statistics, moving averages—everything. This is intentional.

**Why put RNG in Params?** In pure functional programming, randomness is state. If your dropout layer consumes a random key, that's a state change. By threading RNG through `Params`, the signature `outputs, params = model(params, inputs)` tells the whole truth: this function might update some internal state.

This design means every parameter is either:

| Type | Examples | Gradients? | Updated during forward? |
|------|----------|------------|------------------------|
| **Trainable** | weights, embeddings | ✅ Yes | No |
| **Non-trainable** | RNG counters, batch norm stats, EMA | ❌ No | Yes |

The `params.split()` method separates these two categories, which becomes important during training.

## 🎯 Training

During training, we want gradients for trainable parameters but also need to capture updates to non-trainable state (like RNG). The pattern:

```python
@jax.jit(donate_argnames='params')
def train_step(params, inputs, targets):
  # Split into trainable (weights) and non-trainable (RNG, etc).
  trainable, non_trainable = params.split()

  def loss_fn(t, nt):
    # Merge to run the forward pass.
    preds, new_params = model(t.merge(nt), inputs)
    loss = jnp.mean((preds - targets) ** 2)

    # Extract non-trainable parameters updated during the forward pass.
    _, new_nt = new_params.split()
    return loss, new_nt

  # Gradients for trainable, updated state for non-trainable.
  grads, new_non_trainable = jax.grad(loss_fn, has_aux=True)(
      trainable, non_trainable
  )

  # Update the trainable parameters using SGD.
  new_trainable = jax.tree.map(lambda w, g: w - 0.01 * g, trainable, grads)

  # Merge updated trainable and updated non-trainable parameters.
  return new_trainable.merge(new_non_trainable)
```

## 🔀 Batching & Parallel RNG

> **⚠️ JAX Sharp Edge**: This section describes patterns needed due to JAX's PRNG design, not blox design decisions. The main sharp edge is around **initialization with `shard_map`** where different parameters may need different sharding. For `vmap` the patterns are straightforward once understood.

Here is a sharp edge in JAX: if you `vmap` or `shard_map` a function that uses random numbers, every batch element/device gets the *same* random key by default. This means your dropout masks would be identical across the whole batch—defeating the purpose of dropout entirely.

**blox** does not hide this behavior from you. Instead, we give you the tools to handle it explicitly.

### Understanding JAX's Counter-Based PRNG

JAX's PRNG is stateless and deterministic. When you call `rng(params)`, the returned key is computed as:

```python
new_key = jax.random.fold_in(seed, counter)
```

The `seed` is fixed at initialization, and the `counter` increments with each call. This means:

- **Same seed + same counter = same key** (always)
- **Different counter = different key** (even with same seed)

In parallel contexts (`vmap`, `shard_map`), all lanes share the same seed and counter, so they all get identical keys. To get unique randomness per lane, you must "fold in" the lane index.

### The Manual Folding Pattern

The simplest way to understand RNG folding is to pass the batch index explicitly:

```python
graph = bx.Graph('root')
rng = bx.Rng(graph.child('rng'))
dropout = bx.Dropout(graph.child('dropout'), rate=0.5, rng=rng)

def apply_with_explicit_index(params, x, batch_idx):
  # Fold in the batch index to get a unique seed for this lane.
  original_seed = rng.get_seed(params)
  folded_seed = jax.random.fold_in(original_seed, batch_idx)
  params = rng.seed(params, seed=folded_seed)

  out, params = dropout(params, x, is_training=True)

  # Restore original seed (required for replicated params).
  params = rng.seed(params, seed=original_seed)
  return out, params

# Pass jnp.arange(batch_size) as the index.
batch_indices = jnp.arange(4)
outputs, _ = jax.vmap(
    apply_with_explicit_index,
    in_axes=(None, 0, 0),
    out_axes=(0, None),
)(params, batch_inputs, batch_indices)
```

When using `axis_name` with vmap, you can use `jax.lax.axis_index` instead of threading the index through your code. **This is the recommended approach** as it's more idiomatic:

```python
def apply_with_axis_index(params, x):
  original_seed = rng.get_seed(params)
  folded_seed = jax.random.fold_in(
      original_seed, jax.lax.axis_index('batch')
  )
  params = rng.seed(params, seed=folded_seed)

  out, params = dropout(params, x, is_training=True)

  params = rng.seed(params, seed=original_seed)
  return out, params

# axis_name is required for jax.lax.axis_index.
outputs, _ = jax.vmap(
    apply_with_axis_index,
    in_axes=(None, 0),
    out_axes=(0, None),
    axis_name='batch'
)(params, batch_inputs)
```

### Why Restore the Original Seed?

When `params` is replicated across lanes (`out_axes=None`), JAX requires all lanes to return identical pytrees. If each lane has a different folded seed, JAX will error.

Since we're running the same function in each lane, the counter increments by the same amount everywhere. The seed is the only thing that differs (due to folding), so restoring the original seed ensures the params are identical across all lanes.

### Init vs Runtime

During **initialization**, you typically want identical params across all batch elements, so you do NOT fold in the axis index.

During **runtime**, you want unique randomness per batch element (for dropout, etc.), so you DO fold in the axis index.

You can use `params.is_locked` to detect which mode you're in:

```python
def forward(params, x):
  original_seed = rng.get_seed(params)

  # Check if we're in init mode (unlocked) or runtime mode (locked).
  if params.is_locked:
    # Runtime: fold in axis index for unique randomness.
    original_seed = rng.get_seed(params)
    folded_seed = jax.random.fold_in(
        original_seed, jax.lax.axis_index('batch')
    )
    params = rng.seed(params, seed=folded_seed)

  out, params = dropout(params, x, is_training=True)

  # Restore original seed (no-op during init, required for runtime).
  params = rng.seed(params, seed=original_seed)

  return out, params

# Init phase: params are unlocked, no folding.
def init(x):
  params = rng.seed(bx.Params(), seed=42)
  _, params = forward(params, x)
  return params.locked()

# Both init and runtime use the same vmap.
params = jax.vmap(init, axis_name='batch', out_axes=None)(batch_inputs)

# Runtime phase: params are locked, folding is applied.
outputs, _ = jax.vmap(
    forward,
    in_axes=(None, 0),
    axis_name='batch'
)(params, batch_inputs)
```

This pattern lets you use the same `forward` function for both initialization and runtime.

<details>
<summary><strong>Why use JIT instead of shard_map for initialization?</strong></summary>

When initializing models with sharded parameters, **use `jax.jit` with `out_shardings` rather than `shard_map`**.

`shard_map` is tricky for initialization because:
- Different parameters may need different axes folded in (e.g., model axis but not batch axis)
- Multiple model axes mean different params have different sharding requirements
- Managing which axes to fold for which params becomes complex

`jax.jit` is better because:
- Just specify `out_shardings` and JIT handles partitioning automatically
- Use replicated RNG params during init
- JIT is smart about parameter placement during initialization

```python
# RECOMMENDED: Initialize via JIT with out_shardings
@jax.jit(out_shardings=params_sharding)
def init():
    params = rng.seed(bx.Params(), seed=42)
    _, params = model(params, dummy_input)
    return params.locked()

params = init()  # JIT handles sharding automatically
```

</details>

## 📈 Scaling Up

For models that don't fit on one device, you usually need to shard parameters. **blox** lets you specify sharding as metadata when defining layers.

We can use this to initialize parameters directly on the correct devices, avoiding the need to instantiate the full model on the CPU.

```python
from jax.sharding import NamedSharding, PartitionSpec as P

graph = bx.Graph('net')
rng = bx.Rng(graph.child('rng'))

# Define a layer with sharding metadata.
linear = bx.Linear(
    graph.child('linear'),
    output_size=4096,
    rng=rng,
    # Shard weights across the 'model' axis.
    kernel_metadata={'sharding': (None, 'model')},
    bias_metadata={'sharding': ('model',)},
)

def init(x):
  params = rng.seed(bx.Params(), seed=42)
  _, params = linear(params, x)
  return params.locked()

# Get parameter structure without allocating memory or wasting FLOPs.
dummy_input = jnp.ones((4, 4))
abstract_params = jax.eval_shape(init, dummy_input)

# Build sharding specs (assuming 2*2=4 GPU/TPU devices).
mesh = jax.make_mesh((2, 2), ('batch', 'model'))

params_sharding = jax.tree.map(
    lambda p: NamedSharding(mesh, P(*p.sharding)),
    abstract_params,
    is_leaf=lambda x: isinstance(x, bx.Param),
)

 # Example data and output sharding
data_sharding = NamedSharding(mesh, P('batch'))
output_sharding = NamedSharding(mesh, P('batch'))

# Initialize directly on device using out_shardings.
sharded_init = jax.jit(init, out_shardings=params_sharding)
sharded_params = sharded_init(dummy_input)

# Define forward pass with full input/output sharding.
@jax.jit(
    in_shardings=(params_sharding, data_sharding),
    out_shardings=(output_sharding, params_sharding)
)
def forward(params, x):
  return linear(params, x)

out, new_params = forward(sharded_params, dummy_input)
```

## 🔄 Recurrence & Scanning

**blox** provides two base classes for sequence processing:

* **`SequenceBase`**: For models like Transformers where you process the whole sequence at once.
* **`RecurrenceBase`**: For models like LSTMs where you iterate step-by-step.

The built-in `LSTM` and `GRU` extend `RecurrenceBase`. They are designed to work with `jax.lax.scan` for maximum efficiency:

```python
lstm = bx.LSTM(graph.child('lstm'), hidden_size=256, rng=rng)

# Initialize the LSTM state.
inputs = jnp.ones((batch_size, seq_len, features))
state, params = lstm.initial_state(params, inputs[:, 0])

# Process the whole sequence efficiently.
(outputs, final_state), params = lstm.apply(
    params, inputs, prev_state=state
)
```

## 🧠 Under the Hood

**blox** is transparent by design. The abstraction is just automated path handling to keep your code clean and your state pure.

* **Graph**: Defines the model hierarchy (e.g., `net -> mlp -> hidden`). `graph.child('name')` creates child nodes, giving each module a unique path for namespacing its parameters. The graph stores references to all created modules and provides `graph.walk()` for iteration—useful for applying LoRA adapters or toggling training mode across layers.

* **Module**: Has a unique path in the graph and provides convenience methods (`get_param`, `set_param`) to automatically manage its own parameters.

* **Param**: A wrapper around each parameter value that holds a `trainable` flag and arbitrary `metadata`. The trainable flag determines whether a parameter is differentiable or not.

* **Params**: Immutable container holding all state as a flat dictionary of `Param` objects keyed by tuple paths (e.g., `('net', 'mlp', 'hidden', 'kernel')`). Use `split()` to separate trainable from non-trainable state.

* **Rng**: A module that generates deterministic random keys. Since it's used to randomly initialize all other parameters and provide runtime randomness, it must be seeded first via `rng.seed(params, seed=42)`.

## ⚖️ Why the verbosity?

**blox chooses clarity over brevity.**

Most frameworks rely on implicit global state or thread-local contexts to hide parameters. While this saves a few keystrokes, it creates a "black box."

| OOP-style Wrappers | **blox** |
| --- | --- |
| `out = layer(x)` | `outputs, params = layer(params, inputs)` |
| Implicit global state | Explicit state passing |
| Opaque variable scopes | Explicit `bx.Graph` paths |
| Custom `vmap` / `jit` / ... wrappers | Standard `jax.vmap` / `jax.jit` / ... |

By accepting slightly more verbose function signatures, you gain:

1. **Total transparency:** You know exactly what data your function touches.
2. **Full control:** No global state means no unknown side-effects.
3. **Maximum performance:** Zero overhead.

## 🔀 Decoupled Params and Graph

A key design principle is the **clean separation between parameters and the model graph**. Unlike other libraries where params are tightly coupled to modules, blox lets multiple models share the same params.

**Why this separation?**

1. **No single module owns params.** Params are passed in, not stored. Multiple models can use the same params without ownership conflicts.
2. **Avoids pytree complexity.** Modules as pytrees containing both static config and JAX arrays require magic handling, especially for non-hashable types (like lists and dicts). blox keeps a clean split: `Graph` is static (Python objects), `Params` is dynamic (JAX arrays).
3. **Graph is static, params are dynamic.** Graph describes *what* operations to do. Params provide *what values* to use. This separation is maintained throughout execution.

**Use cases:**

* **Actor vs Learner in RL**: Separate models for data collection and training that share weights.
* **Training vs Evaluation**: Scenarios where evaluation logic differs significantly while relying on the same parameters.

**Rule of thumb:** If the model needs recompilation, prefer creating a new model instead of modifying an existing one or creating an Uber-module.

This design should be inspired by the function. For example, if you need both static and dynamic scanning behavior, you should opt for creating two models.

```python
# Create two LSTM variants with the same parameter structure.
def create_lstm(is_static: bool):
  graph = bx.Graph('model')
  rng = bx.Rng(graph.child('rng'))
  return bx.LSTM(
      graph.child('lstm'),
      hidden_size=64,
      rng=rng,
      is_static=is_static
  )

lstm_static = create_lstm(is_static=True)  # Python loop (debuggable).
lstm_dynamic = create_lstm(is_static=False)  # lax.scan (production).

# Both models work with the SAME params!
out_static, _ = lstm_static.apply(params, inputs, prev_state=state)
out_dynamic, _ = lstm_dynamic.apply(params, inputs, prev_state=state)
```

While nothing prevents users from changing modules in place, JAX will not recompile functions automatically unless manually instructed (see [JAX Gotchas](https://docs.jax.dev/en/latest/notebooks/Common_Gotchas_in_JAX.html#using-jax-jit-with-class-methods)). We do support exceptions for ergonomics, such as the `is_training` when evaluation logic is only a simple flag away from training logic, e.g. dropout.

## 📄 License

MIT. See [LICENSE](LICENSE).