hanzo_ml/
lib.rs

1//! ML framework for Rust
2//!
3//! ```rust
4//! use hanzo_ml::{Tensor, DType, Device};
5//! # use hanzo_ml::Error;
6//! # fn main() -> Result<(), Error>{
7//!
8//! let a = Tensor::arange(0f32, 6f32, &Device::Cpu)?.reshape((2, 3))?;
9//! let b = Tensor::arange(0f32, 12f32, &Device::Cpu)?.reshape((3, 4))?;
10//! let c = a.matmul(&b)?;
11//!
12//! # Ok(())}
13//! ```
14//!
15//! ## Features
16//!
17//! - Simple syntax (looks and feels like PyTorch)
18//! - CPU and Cuda backends (and M1 support)
19//! - Enable serverless (CPU) small and fast deployments
20//! - Model training
21//! - Distributed computing (NCCL).
22//! - Models out of the box (Llama, Whisper, Falcon, ...)
23//!
24//! ## FAQ
25//!
26//! - Why Hanzo?
27//!
28//! Hanzo stems from the need to reduce binary size in order to *enable serverless*
29//! possible by making the whole engine smaller than PyTorch very large library volume
30//!
31//! And simply *removing Python* from production workloads.
32//! Python can really add overhead in more complex workflows and the [GIL](https://www.backblaze.com/blog/the-python-gil-past-present-and-future/) is a notorious source of headaches.
33//!
34//! Rust is cool, and a lot of the HF ecosystem already has Rust crates [safetensors](https://github.com/huggingface/safetensors) and [tokenizers](https://github.com/huggingface/tokenizers)
35//!
36//! ## Other Crates
37//!
38//! Hanzo consists of a number of crates. This crate holds core the common data structures but you may wish
39//! to look at the docs for the other crates which can be found here:
40//!
41//! - [hanzo-ml](https://docs.rs/hanzo-ml/). Core Datastructures and DataTypes.
42//! - [hanzo-nn](https://docs.rs/hanzo-nn/). Building blocks for Neural Nets.
43//! - [hanzo-datasets](https://docs.rs/hanzo-datasets/). Rust access to commonly used Datasets like MNIST.
44//! - [hanzo-ml-examples](https://docs.rs/hanzo-ml-examples/). Examples of Hanzo in Use.
45//! - [hanzo-onnx](https://docs.rs/hanzo-onnx/). Loading and using ONNX models.
46//! - [hanzo-ml-pyo3](https://docs.rs/hanzo-ml-pyo3/). Access to Hanzo from Python.
47//! - [hanzo-transformers](https://docs.rs/hanzo-transformers/). Hanzo implementation of many published transformer models.
48//!
49
50#[cfg(feature = "accelerate")]
51mod accelerate;
52pub mod backend;
53pub mod backprop;
54pub mod compute;
55pub mod conv;
56mod convert;
57pub mod cpu;
58pub mod cpu_backend;
59#[cfg(feature = "cuda")]
60pub mod cuda_backend;
61mod custom_op;
62mod device;
63pub mod display;
64mod dtype;
65pub mod dummy_cuda_backend;
66pub mod dummy_dtype;
67mod dummy_metal_backend;
68pub mod dummy_vulkan_backend;
69pub mod dummy_wgpu_backend;
70pub mod error;
71pub mod fattn_decode;
72mod hc_sinkhorn;
73mod indexer;
74pub mod layout;
75#[cfg(feature = "metal")]
76pub mod metal_backend;
77#[cfg(feature = "mkl")]
78mod mkl;
79pub mod model_delta;
80pub mod nditer;
81pub mod npy;
82pub mod op;
83pub mod pickle;
84pub mod quantized;
85#[cfg(feature = "rocm")]
86pub mod rocm_backend;
87pub mod safetensors;
88pub mod scalar;
89pub mod shape;
90mod sort;
91mod storage;
92pub mod streaming;
93mod strided_index;
94mod tensor;
95mod tensor_cat;
96pub mod test_utils;
97pub mod utils;
98mod variable;
99#[cfg(feature = "vulkan")]
100pub mod vulkan_backend;
101#[cfg(feature = "wgpu")]
102pub mod wgpu_backend;
103
104#[cfg(feature = "cudnn")]
105pub use cuda_backend::cudnn;
106
107pub use cpu_backend::{CpuStorage, CpuStorageRef};
108#[cfg(feature = "ug")]
109pub use custom_op::UgIOp1;
110pub use custom_op::{CustomOp1, CustomOp2, CustomOp3, InplaceOp1, InplaceOp2, InplaceOp3};
111pub use device::{Device, DeviceLocation, NdArray};
112pub use dtype::{DType, DTypeParseError, FloatDType, IntDType, WithDType};
113pub use dummy_dtype::{F4, F6E2M3, F6E3M2, F8E8M0};
114pub use error::{Context, Error, Result};
115pub use indexer::{IndexOp, TensorIndexer};
116pub use layout::Layout;
117pub use nditer::NdIter;
118pub use shape::{Shape, D};
119pub use storage::Storage;
120pub use streaming::{StreamTensor, StreamingBinOp, StreamingModule};
121pub use strided_index::{StridedBlocks, StridedIndex};
122pub use tensor::{Tensor, TensorId};
123pub use variable::Var;
124
125#[cfg(feature = "cuda")]
126pub use cuda_backend as cuda;
127
128#[cfg(not(feature = "cuda"))]
129pub use dummy_cuda_backend as cuda;
130
131pub use cuda::{CudaDevice, CudaStorage};
132
133#[cfg(feature = "metal")]
134pub use metal_backend::{MetalDevice, MetalError, MetalStorage};
135
136#[cfg(not(feature = "metal"))]
137pub use dummy_metal_backend::{MetalDevice, MetalError, MetalStorage};
138
139#[cfg(feature = "rocm")]
140pub use rocm_backend::{
141    set_force_scalar_matvec, RocmDevice, RocmError, RocmQuantType, RocmStorage,
142};
143
144#[cfg(feature = "vulkan")]
145pub use vulkan_backend as vulkan;
146
147#[cfg(not(feature = "vulkan"))]
148pub use dummy_vulkan_backend as vulkan;
149
150pub use vulkan::{VulkanDevice, VulkanStorage};
151
152#[cfg(feature = "vulkan")]
153pub use vulkan::{PagedAttnArgs, ReshapeCacheArgs, VkGraph, VkGraphAttn};
154
155// `wgpu_be` alias to avoid clashing with the `wgpu` crate name.
156#[cfg(feature = "wgpu")]
157pub use wgpu_backend as wgpu_be;
158
159#[cfg(not(feature = "wgpu"))]
160pub use dummy_wgpu_backend as wgpu_be;
161
162pub use wgpu_be::{WgpuDevice, WgpuStorage};
163
164#[cfg(feature = "mkl")]
165extern crate intel_mkl_src;
166
167#[cfg(feature = "accelerate")]
168extern crate accelerate_src;
169
170pub trait ToUsize2 {
171    fn to_usize2(self) -> (usize, usize);
172}
173
174impl ToUsize2 for usize {
175    fn to_usize2(self) -> (usize, usize) {
176        (self, self)
177    }
178}
179
180impl ToUsize2 for (usize, usize) {
181    fn to_usize2(self) -> (usize, usize) {
182        self
183    }
184}
185
186/// Defining a module with forward method using a single argument.
187pub trait Module {
188    fn forward(&self, xs: &Tensor) -> Result<Tensor>;
189}
190
191impl<T: Fn(&Tensor) -> Result<Tensor>> Module for T {
192    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
193        self(xs)
194    }
195}
196
197impl<M: Module> Module for Option<&M> {
198    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
199        match self {
200            None => Ok(xs.clone()),
201            Some(m) => m.forward(xs),
202        }
203    }
204}
205
206/// A single forward method using a single single tensor argument and a flag to
207/// separate the training and evaluation behaviors.
208pub trait ModuleT {
209    fn forward_t(&self, xs: &Tensor, train: bool) -> Result<Tensor>;
210}
211
212impl<M: Module> ModuleT for M {
213    fn forward_t(&self, xs: &Tensor, _train: bool) -> Result<Tensor> {
214        self.forward(xs)
215    }
216}
hanzo_ml/lib.rs

hanzo_ml/
lib.rs
