Metadata-Version: 2.5
Name: ion-nn
Version: 0.16.0
Summary: A simple library for neural and graph networks in JAX.
Project-URL: Homepage, https://github.com/auxeno/ion
Project-URL: Documentation, https://auxeno.github.io/ion/
Project-URL: Repository, https://github.com/auxeno/ion
Author: Alex Goddard
License-Expression: Apache-2.0
License-File: LICENSE
Keywords: deep-learning,jax,neural-networks
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.11
Requires-Dist: jax>=0.7.2
Requires-Dist: jaxtyping>=0.3.0
Requires-Dist: optax
Requires-Dist: treescope
Description-Content-Type: text/markdown

<div align="center">

  <h1><img src="https://raw.githubusercontent.com/auxeno/ion/main/assets/logo-transparent.png" alt="Ion" width="72"><br>Ion</h1>

  <h3>A simple library for neural and graph networks in JAX.</h3>

[![Python](https://img.shields.io/badge/Python-3.11+-7C3AED.svg)](https://www.python.org/)
[![PyPI](https://img.shields.io/pypi/v/ion-nn?color=478AF5)](https://pypi.org/project/ion-nn/)
[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json&color=313131&labelColor=555555)](https://github.com/astral-sh/ruff)
[![CI](https://github.com/auxeno/ion/actions/workflows/ci.yml/badge.svg)](https://github.com/auxeno/ion/actions/workflows/ci.yml)
[![codecov](https://codecov.io/gh/auxeno/ion/graph/badge.svg)](https://codecov.io/gh/auxeno/ion)

</div>

---

Ion is a simple neural network library for JAX. The core introduces four concepts
(`Module`, `Param`, `Buffer`, `Optimizer`) that make it simple to build and train
neural networks. Models are [pytrees](https://docs.jax.dev/en/latest/pytrees.html)
that work directly with native JAX transforms. Ion also ships neural and graph
network layers built on the core.

```bash
pip install ion-nn
```

The [documentation](https://auxeno.github.io/ion/) covers the core, layers, and common workflows.

## Example

A model built from Ion's standard layers, trained with native JAX transforms:

```python
import jax, optax, typing

import ion
import ion.nn as nn


class MLP(nn.Module):
    layer_1: nn.Linear
    layer_2: nn.Linear
    activation: typing.Callable

    def __init__(self, activation=jax.nn.relu, *, key):
        keys = jax.random.split(key, 2)
        self.layer_1 = nn.Linear(784, 128, key=keys[0])
        self.layer_2 = nn.Linear(128, 10, key=keys[1])
        self.activation = activation

    def __call__(self, x):
        return self.layer_2(self.activation(self.layer_1(x)))


def loss_fn(model, x, y):
    logits = model(x)
    return optax.softmax_cross_entropy_with_integer_labels(logits, y).mean()


@jax.jit
def train_step(model, optimizer, x, y):
    grads = jax.grad(loss_fn)(model, x, y)
    model, optimizer = optimizer.update(model, grads)
    return model, optimizer


model = MLP(key=jax.random.key(0))

optimizer = ion.Optimizer(optax.adam(3e-4), model)

for x, y in data:
    model, optimizer = train_step(model, optimizer, x, y)
```

## Documentation

- [Overview](https://auxeno.github.io/ion/overview/) - the core abstractions and design
- [Core](https://auxeno.github.io/ion/core/module/) - `Module`, `Param`, `Buffer`, and `Optimizer`
- [NN guide](https://auxeno.github.io/ion/nn/guide/) and [GNN guide](https://auxeno.github.io/ion/gnn/guide/) - array formats and shared conventions
- [Workflows](https://auxeno.github.io/ion/workflows/) - freezing, mixed precision, serialization, inspecting models
- [Sharp edges](https://auxeno.github.io/ion/sharp-edges/) - known limitations and gotchas
- [Examples](https://auxeno.github.io/ion/examples/) - end-to-end training scripts and notebooks
- [Benchmarks](https://auxeno.github.io/ion/benchmarks/) - comparisons with Equinox, Flax NNX, and PyTorch

## Layers

Ion ships with standard neural network layers. Each is a `Module` with trainable
`Param` leaves; stateful layers also hold non-trainable `Buffer` fields.

| Category        | Layers                                                                    |
|-----------------|---------------------------------------------------------------------------|
| Linear          | [`Linear`](https://auxeno.github.io/ion/nn/layers/linear/#ion.nn.Linear), [`Identity`](https://auxeno.github.io/ion/nn/layers/linear/#ion.nn.Identity) |
| Convolution     | [`Conv`](https://auxeno.github.io/ion/nn/layers/conv/#ion.nn.Conv), [`ConvTranspose`](https://auxeno.github.io/ion/nn/layers/conv/#ion.nn.ConvTranspose) |
| Attention       | [`MultiHeadAttention`](https://auxeno.github.io/ion/nn/layers/attention/#ion.nn.MultiHeadAttention) |
| Normalization   | [`LayerNorm`](https://auxeno.github.io/ion/nn/layers/norm/#ion.nn.LayerNorm), [`RMSNorm`](https://auxeno.github.io/ion/nn/layers/norm/#ion.nn.RMSNorm), [`BatchNorm`](https://auxeno.github.io/ion/nn/layers/norm/#ion.nn.BatchNorm), [`GroupNorm`](https://auxeno.github.io/ion/nn/layers/norm/#ion.nn.GroupNorm), [`SpectralNorm`](https://auxeno.github.io/ion/nn/layers/norm/#ion.nn.SpectralNorm) |
| Recurrent       | [`RNN`](https://auxeno.github.io/ion/nn/layers/recurrent/#ion.nn.RNN), [`LSTM`](https://auxeno.github.io/ion/nn/layers/recurrent/#ion.nn.LSTM), [`GRU`](https://auxeno.github.io/ion/nn/layers/recurrent/#ion.nn.GRU) |
| SSM             | [`S4D`](https://auxeno.github.io/ion/nn/layers/ssm/#ion.nn.S4D), [`S5`](https://auxeno.github.io/ion/nn/layers/ssm/#ion.nn.S5) |
| Pooling         | [`MaxPool`](https://auxeno.github.io/ion/nn/layers/pool/#ion.nn.MaxPool), [`AvgPool`](https://auxeno.github.io/ion/nn/layers/pool/#ion.nn.AvgPool) |
| Embedding       | [`Embedding`](https://auxeno.github.io/ion/nn/layers/embedding/#ion.nn.Embedding) |
| Positional      | [`RoPE`](https://auxeno.github.io/ion/nn/layers/positional/#ion.nn.RoPE), [`LearnedPositionalEmbedding`](https://auxeno.github.io/ion/nn/layers/positional/#ion.nn.LearnedPositionalEmbedding), [`SinusoidalPositionalEmbedding`](https://auxeno.github.io/ion/nn/layers/positional/#ion.nn.SinusoidalPositionalEmbedding) |
| Stochastic      | [`Dropout`](https://auxeno.github.io/ion/nn/layers/stochastic/#ion.nn.Dropout) |
| Blocks          | [`Sequential`](https://auxeno.github.io/ion/nn/layers/sequential/#ion.nn.Sequential), [`MLP`](https://auxeno.github.io/ion/nn/layers/mlp/#ion.nn.MLP) |

Graph layers live in `ion.gnn` and take node features with COO `senders`/`receivers`
edge indices.

| Category        | Layers                                                                    |
|-----------------|---------------------------------------------------------------------------|
| Convolution     | [`GCNConv`](https://auxeno.github.io/ion/gnn/layers/conv/#ion.gnn.GCNConv), [`GraphConv`](https://auxeno.github.io/ion/gnn/layers/conv/#ion.gnn.GraphConv), [`SAGEConv`](https://auxeno.github.io/ion/gnn/layers/conv/#ion.gnn.SAGEConv) |
| Attention       | [`GATConv`](https://auxeno.github.io/ion/gnn/layers/attention/#ion.gnn.GATConv), [`GATv2Conv`](https://auxeno.github.io/ion/gnn/layers/attention/#ion.gnn.GATv2Conv), [`TransformerConv`](https://auxeno.github.io/ion/gnn/layers/attention/#ion.gnn.TransformerConv) |
| Isomorphism     | [`GINConv`](https://auxeno.github.io/ion/gnn/layers/isomorphism/#ion.gnn.GINConv), [`GINEConv`](https://auxeno.github.io/ion/gnn/layers/isomorphism/#ion.gnn.GINEConv) |
| Composite       | [`GraphNetwork`](https://auxeno.github.io/ion/gnn/layers/composite/#ion.gnn.GraphNetwork), [`EdgeUpdate`](https://auxeno.github.io/ion/gnn/layers/composite/#ion.gnn.EdgeUpdate), [`NodeUpdate`](https://auxeno.github.io/ion/gnn/layers/composite/#ion.gnn.NodeUpdate) |
| Relational      | [`RGCNConv`](https://auxeno.github.io/ion/gnn/layers/relational/#ion.gnn.RGCNConv), [`HGTConv`](https://auxeno.github.io/ion/gnn/layers/relational/#ion.gnn.HGTConv) |
| Gated           | [`GatedGCNConv`](https://auxeno.github.io/ion/gnn/layers/gated/#ion.gnn.GatedGCNConv) |
| Pooling         | [`GlobalAttentionPool`](https://auxeno.github.io/ion/gnn/layers/pool/#ion.gnn.GlobalAttentionPool), [`MultiHeadAttentionPool`](https://auxeno.github.io/ion/gnn/layers/pool/#ion.gnn.MultiHeadAttentionPool) |

See the [NN guide](https://auxeno.github.io/ion/nn/guide/) and [GNN guide](https://auxeno.github.io/ion/gnn/guide/) for array formats, spatial layers, and shared conventions.

## License

Released under the Apache License 2.0.

## Citation

To cite this repository:

```bibtex
@software{ion,
  title = {Ion: A simple library for neural and graph networks in JAX.},
  author = {Alex Goddard},
  url = {https://github.com/auxeno/ion},
  year = {2026}
}
```
</content>
</invoke>
