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 more
• Deep integration with (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 ** 2, axis=0)

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

Roadmap

• ✅ Function Transformations: vmap, jit, vjp, jvp, grad
• ✅ Mojo Kernel Integration: CPU/GPU acceleration working
• 👷 Extended Operations: Comprehensive math operations
• 💡 Enhanced Mojo API: When Mojo/MAX ecosystem stabilizes

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/                  # Function transformations and core array class
│   ├── ops/                   # Mathematical operations (binary, unary, linalg, etc.)
│   ├── nn/                    # Neural network modules and models
│   └── utils/                 # Utilities (broadcasting, formatting, types)
├── tests/                     # Comprehensive test suite
├── examples/                  # Towards MLP training and other usage examples
└── experimental/              # An alternative pure Mojo API (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!