vatra-py 0.2.1
A lightweight Deep Learning framework backed by Vulkan (wgpu)
License: MIT
Keywords: deep-learning, rust, vulkan, pytorch, autograd

# Vatra — Vulkan Neural Network Library

**Vatra** (Ukrainian for "bonfire" 🔥) is a high-performance neural network library written in Rust, powered by Vulkan (via `wgpu`). It combines the ergonomic API of **PyTorch** with the functional semantics and automatic differentiation of **JAX**.

## Features

- **🔥 GPU Acceleration**: Runs on any GPU supporting Vulkan/Metal/DX12 (via `wgpu`).
- **🧠 PyTorch-like API**: Familiar `Tensor`, `Linear`, `Conv2d`, `Sequential` abstractions.
- **⚡ JAX-like Autograd**: Functional transformations like `grad()`, `vmap()`, and `scan()`.
- **🐍 Python Bindings**: Use Vatra directly from Python.
- **🎯 Auto Device Selection**: `Device::auto()` picks the best available device.

## Quick Start

```rust
use vatra::prelude::*;

fn main() -> VatraResult<()> {
    // Automatic device selection (GPU if available, else CPU)
    set_default_device(Device::auto());

    // Create tensors on the best available device
    let x = Tensor::randn_d([32, 784])?;

    // Define model
    let device = get_default_device();
    let model = Sequential::new()
        .add(Linear::new(784, 256, true, device)?)
        .add(ReLU)
        .add(Linear::new(256, 10, true, device)?);

    // Forward pass
    let output = model.forward(&x)?;
    println!("Output shape: {}", output.shape());

    Ok(())
}
```

## Python Usage

Install via `pip`:

```bash
pip install vatra
```

Or build from source:

```bash
cd crates/vatra-py
maturin develop --release
```

Example:

```python
import vatra

# Auto device selection (like PyTorch)
print(f"GPU available: {vatra.cuda_is_available()}")
vatra.set_default_device("auto")

# Create tensors (uses default device)
x = vatra.Tensor.randn([32, 784])
layer = vatra.Linear(784, 256)

# Forward pass
y = layer.forward(x)
print(f"Output shape: {y.shape}")

# Autograd
def loss_fn(x):
    y = layer.forward(x)
    return vatra.sum(y)

grad_fn = vatra.value_and_grad(loss_fn)
val, grads = grad_fn(x)
print(f"Loss: {val.item()}")
```

## Device Selection

```rust
// Check GPU availability
if Device::cuda_is_available() {
    println!("Found {} GPU(s)", Device::device_count());
}

// Automatic selection
let device = Device::auto();

// Or explicit
let gpu = Device::Vulkan(0);
let cpu = Device::Cpu;

// Set global default
set_default_device(Device::auto());

// Use _d suffix methods for default device
let x = Tensor::randn_d([100, 100])?;
let y = Tensor::zeros_d([100, 100])?;
```

## Architecture

| Crate            | Purpose                         |
| ---------------- | ------------------------------- |
| `vatra-gpu`      | wgpu backend, WGSL shaders      |
| `vatra-core`     | Tensor, Device, Shape, DType    |
| `vatra-ops`      | GPU kernels for math operations |
| `vatra-autograd` | Reverse-mode AD, grad(), vmap() |
| `vatra-nn`       | Neural network layers           |
| `vatra-optim`    | Optimizers (SGD, Adam, AdamW)   |
| `vatra-py`       | Python bindings (PyO3)          |
| `vatra`          | Facade crate with prelude       |

## Documentation

See the [docs](./docs) folder for:

- [Installation Guide](./docs/installation.md)
- [Quick Start](./docs/quickstart.md)
- [API Reference](./docs/api/README.md)
- [Tutorials](./docs/tutorials/README.md)

## License

MIT