Skip to main content

hanzo_nn/
lib.rs

1//! hanzo-nn
2//!
3//! ## Other Crates
4//!
5//! Hanzo consists of a number of crates. This crate holds structs and functions
6//! that allow you to build and train neural nets. You may wish
7//! to look at the docs for the other crates which can be found here:
8//!
9//! - [hanzo-ml](https://docs.rs/hanzo-ml/). Core Datastructures and DataTypes.
10//! - [hanzo-nn](https://docs.rs/hanzo-nn/). Building blocks for Neural Nets.
11//! - [hanzo-datasets](https://docs.rs/hanzo-datasets/). Rust access to commonly used Datasets like MNIST.
12//! - [hanzo-ml-examples](https://docs.rs/hanzo-ml-examples/). Examples of Hanzo in Use.
13//! - [hanzo-onnx](https://docs.rs/hanzo-onnx/). Loading and using ONNX models.
14//! - [hanzo-ml-pyo3](https://docs.rs/hanzo-ml-pyo3/). Access to Hanzo from Python.
15//! - [hanzo-transformers](https://docs.rs/hanzo-transformers/). Hanzo implementation of many published transformer models.
16//!
17
18pub mod activation;
19pub mod attention;
20pub mod batch_norm;
21pub mod conv;
22pub mod cpu_flash_attention;
23pub mod embedding;
24pub mod encoding;
25pub mod func;
26pub mod group_norm;
27pub mod init;
28pub mod kv_cache;
29pub mod layer_norm;
30pub mod linear;
31pub mod loss;
32pub mod moe;
33pub mod ops;
34pub mod optim;
35pub mod rnn;
36pub mod rotary_emb;
37pub mod sampling;
38pub mod sequential;
39pub mod var_builder;
40pub mod var_map;
41/// Re-export of [`attention::varlen`] for backward compatibility.
42pub mod varlen_attention {
43    pub use crate::attention::varlen::*;
44}
45
46pub use activation::{prelu, Activation, PReLU};
47pub use batch_norm::{batch_norm, BatchNorm, BatchNormConfig};
48pub use conv::{
49    conv1d, conv1d_no_bias, conv2d, conv2d_no_bias, conv_transpose1d, conv_transpose1d_no_bias,
50    conv_transpose2d, conv_transpose2d_no_bias, Conv1d, Conv1dConfig, Conv2d, Conv2dConfig,
51    ConvTranspose1d, ConvTranspose1dConfig, ConvTranspose2d, ConvTranspose2dConfig,
52};
53pub use embedding::{embedding, Embedding};
54pub use func::{func, func_t, Func, FuncT};
55pub use group_norm::{group_norm, GroupNorm};
56pub use init::Init;
57pub use layer_norm::{
58    layer_norm, layer_norm_no_bias, rms_norm, LayerNorm, LayerNormConfig, RmsNorm,
59};
60pub use linear::{linear, linear_b, linear_no_bias, Linear};
61pub use ops::Dropout;
62pub use optim::{AdamW, Optimizer, ParamsAdamW, SGD};
63pub use rnn::{gru, lstm, GRUConfig, LSTMConfig, GRU, LSTM, RNN};
64pub use sequential::{seq, Sequential};
65pub use var_builder::VarBuilder;
66pub use var_map::VarMap;
67
68pub use hanzo_ml::{Module, ModuleT};