NABLA

NABLA#

Nabla is a Python library that provides three key features:

Multidimensional Array computation (like NumPy) with strong GPU acceleration
Composable Function Transformations: vmap, grad, jit, and other Automatic Differentiation tools
Deep integration with MAX and (custom) Mojo kernels

Installation#

📦 Now available on PyPI!

pip install nabla-ml

Note: Nabla also includes an experimental pure Mojo API for native Mojo development.

Quick Start#

import nabla as nb

# Example function using Nabla's array operations
def foo(input):
    return nb.sum(input * input, axes=0)

# Vectorize, differentiate, accelerate
foo_grads = nb.jit(nb.vmap(nb.grad(foo)))
gradients = foo_grads(nb.randn((10, 5)))

Development Setup#

For contributors and advanced users:

# Clone and install in development mode
git clone https://github.com/nabla-ml/nb.git
cd nabla
pip install -e ".[dev]"

# Run tests
pytest

# Format and lint code
ruff format nabla/
ruff check nabla/ --fix

Repository Structure#

nabla/
├── nabla/                     # Core Python library
│   ├── core/                  # Array class and MAX compiler integration
│   ├── nn/                    # Neural network modules and models
│   ├── ops/                   # Mathematical operations (binary, unary, linalg, etc.)
│   ├── transforms/            # Function transformations (vmap, grad, jit, etc.)
│   └── utils/                 # Utilities (formatting, types, MAX-interop, etc.)
├── tests/                     # Comprehensive test suite
├── tutorials/                 # Notebooks on Nabla usage for ML tasks
├── examples/                  # Example scripts for common use cases
└── experimental/              # Core (pure) Mojo library (WIP!)

Contributing#

Contributions welcome! Discuss significant changes in Issues first. Submit PRs for bugs, docs, and smaller features.

License#

Nabla is licensed under the Apache-2.0 license.

Thank you for checking out Nabla!