rai-core 0.6.0

ML framework with Ergonomic APIs in Rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79

# RAI

![Rust](https://github.com/cksac/rai/workflows/Rust/badge.svg)
[![Docs Status](https://docs.rs/rai/badge.svg)](https://docs.rs/rai)
[![Latest Version](https://img.shields.io/crates/v/rai.svg)](https://crates.io/crates/rai)

ML framework with Ergonomic APIs in Rust. Lazy computation and composable transformations.
---

## Installation
```sh
cargo add rai
```

## Code snippets
### Function transformations (jvp, vjp, grad, value_and_grad)
```rust
use rai::{backend::Cpu, grad, DType, Func, Tensor};

fn f(x: &Tensor) -> Tensor {
    x.sin()
}

fn main() {
    let grad_fn = grad(grad(f));

    let backend = &Cpu;
    let x = Tensor::ones([1], DType::F32, backend);
    let grad = grad_fn.apply((&x,));

    println!("{}", grad.dot_graph());
    println!("{}", grad);
}
```

### NN Modules, Optimizer and loss functions
```rust
fn loss_fn<M: DifferentiableModule + 'static>(
    model: &M,
    input: &Tensor,
    labels: &Tensor,
) -> (Tensor, Aux<Tensor>) {
    let logits = model.forward(input);
    let loss = softmax_cross_entropy(&logits, labels).mean(..);
    (loss, Aux(logits))
}

fn train_step<O: Optimizer, M: DifferentiableModule + 'static>(
    optimizer: &mut O,
    model: &M,
    input: &Tensor,
    labels: &Tensor,
) {
    let vg_fn = value_and_grad(loss_fn);
    let ((_loss, Aux(_logits)), (grads, ..)) = vg_fn.apply((model, input, labels));
    let mut params = optimizer.step(&grads);
    eval(&params);
    model.update(&mut params);
}
```

## Examples
- [linear_regression](https://github.com/cksac/rai/blob/main/examples/linear_regression/src/main.rs)
    - `cargo run --bin linear_regression --release`
- [mnist](https://github.com/cksac/rai/blob/main/examples/mnist/src/main.rs)
    - `cargo run --bin mnist --release`



# LICENSE

This project is licensed under either of

- Apache License, Version 2.0, ([LICENSE-APACHE](LICENSE-APACHE) or
  http://www.apache.org/licenses/LICENSE-2.0)
- MIT license ([LICENSE-MIT](LICENSE-MIT) or
  http://opensource.org/licenses/MIT)

at your option.