Metadata-Version: 2.4
Name: mlx-predictive-coding
Version: 0.1.0
Summary: A pure, decoupled Predictive Coding framework for Apple MLX.
Author-email: XY <dev@yuuphoria.com>
Classifier: Programming Language :: Python :: 3
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Requires-Python: >=3.9
Description-Content-Type: text/markdown
Requires-Dist: mlx>=0.10.0

# mlx-pc

**A straightforward, minimal Predictive Coding (PC) implementation for the Apple MLX framework.**

`mlx-pc` provides stateless predictive coding layers and a network container to manage relaxation dynamics. It is designed to be completely decoupled from the optimizer, allowing you to plug in standard MLX optimizers or custom implementations.

## Installation

```bash
pip install mlx-pc
```

## API Reference
The library exposes two main components: PCNetwork and PCLayer.

1. PCNetwork (Container)
The primary module that stacks PCLayers and handles the temporal state management (relaxation iterations) during the forward pass.

* Initialization:
```bash
model = PCNetwork(layer_dims: list[int], bias=True)
```
layer_dims: A list of integers defining the feature dimensions of each layer (e.g., [128, 256, 512]).

* Forward Pass (__call__):
```bash
predictions, layer_errors = model(sensory_x, max_iters=10, base_eta=0.1, alpha=1.0, beta=0.5)
```
x: Input tensor of shape (batch_size, seq_len, layer_dims[0]).
max_iters: Number of internal relaxation steps to perform for the fast variables.
Returns: A tuple containing the final predictions and a list of layer_errors tensors.

2. PCLayer (Stateless Block)
A single predictive coding layer. It acts as a pure function, computing local surprisal and the updated state based on bottom-up inputs and top-down priors.

* Initialization:
```bash
layer = PCLayer(in_dim: int, out_dim: int, bias=True)
```

* Forward Pass (__call__):
```bash
current_state, error = layer(x, higher_state, prev_state, base_eta=0.1, alpha=1.0, beta=0.5)
```
x: The bottom-up input from the lower layer.
higher_state: The expected state from the higher layer.
prev_state: The state of this layer at iteration t-1.
Returns: The updated current_state and the calculated error (residual).

## Quick Integration Example

mlx-pc relies on standard MLX auto-grad to update weights. You extract the total squared error (Free Energy) and pass it to any optimizer.

```bash
import mlx.core as mx
import mlx.optimizers as optim
from mlx_pc import PCNetwork

# 1. Init
model = PCNetwork(layer_dims=[128, 256, 512])
optimizer = optim.AdamW(learning_rate=1e-3)

# 2. Define Loss
def loss_fn(model, x):
    predictions, layer_errors = model(x)
    return mx.sum(mx.array([mx.sum(mx.square(e)) for e in layer_errors]))

loss_and_grad_fn = mx.value_and_grad(model, loss_fn)

# 3. Step
x = mx.random.normal((4, 32, 128))
loss, grads = loss_and_grad_fn(model, x)
optimizer.update(model, grads)
mx.eval(model.parameters(), optimizer.state)
```
