candle_core/
lib.rs

1//! ML framework for Rust
2//!
3//! ```rust
4//! use candle_core::{Tensor, DType, Device};
5//! # use candle_core::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 Candle?
27//!
28//! Candle 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//! Candle 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//! - [candle-core](https://docs.rs/candle-core/). Core Datastructures and DataTypes.
42//! - [candle-nn](https://docs.rs/candle-nn/). Building blocks for Neural Nets.
43//! - [candle-datasets](https://docs.rs/candle-datasets/). Rust access to commonly used Datasets like MNIST.
44//! - [candle-examples](https://docs.rs/candle-examples/). Examples of Candle in Use.
45//! - [candle-onnx](https://docs.rs/candle-onnx/). Loading and using ONNX models.
46//! - [candle-pyo3](https://docs.rs/candle-pyo3/). Access to Candle from Python.
47//! - [candle-transformers](https://docs.rs/candle-transformers/). Candle implemntation of many published transformer models.
48//!
49
50#[cfg(feature = "accelerate")]
51mod accelerate;
52pub mod backend;
53pub mod backprop;
54pub mod conv;
55mod convert;
56pub mod cpu;
57pub mod cpu_backend;
58#[cfg(feature = "cuda")]
59pub mod cuda_backend;
60mod custom_op;
61mod device;
62pub mod display;
63mod dtype;
64pub mod dummy_cuda_backend;
65mod dummy_metal_backend;
66pub mod error;
67mod indexer;
68pub mod layout;
69#[cfg(feature = "metal")]
70pub mod metal_backend;
71#[cfg(feature = "mkl")]
72mod mkl;
73pub mod npy;
74pub mod op;
75pub mod pickle;
76pub mod quantized;
77pub mod safetensors;
78pub mod scalar;
79pub mod shape;
80mod sort;
81mod storage;
82pub mod streaming;
83mod strided_index;
84mod tensor;
85mod tensor_cat;
86pub mod test_utils;
87pub mod utils;
88mod variable;
89
90#[cfg(feature = "cudnn")]
91pub use cuda_backend::cudnn;
92
93pub use cpu_backend::{CpuStorage, CpuStorageRef};
94pub use custom_op::{CustomOp1, CustomOp2, CustomOp3, InplaceOp1, InplaceOp2, InplaceOp3, UgIOp1};
95pub use device::{Device, DeviceLocation, NdArray};
96pub use dtype::{DType, DTypeParseError, FloatDType, IntDType, WithDType};
97pub use error::{Context, Error, Result};
98pub use indexer::{IndexOp, TensorIndexer};
99pub use layout::Layout;
100pub use shape::{Shape, D};
101pub use storage::Storage;
102pub use streaming::{StreamTensor, StreamingBinOp, StreamingModule};
103pub use strided_index::{StridedBlocks, StridedIndex};
104pub use tensor::{Tensor, TensorId};
105pub use variable::Var;
106
107#[cfg(feature = "cuda")]
108pub use cuda_backend as cuda;
109
110#[cfg(not(feature = "cuda"))]
111pub use dummy_cuda_backend as cuda;
112
113pub use cuda::{CudaDevice, CudaStorage};
114
115#[cfg(feature = "metal")]
116pub use metal_backend::{MetalDevice, MetalError, MetalStorage};
117
118#[cfg(not(feature = "metal"))]
119pub use dummy_metal_backend::{MetalDevice, MetalError, MetalStorage};
120
121#[cfg(feature = "mkl")]
122extern crate intel_mkl_src;
123
124#[cfg(feature = "accelerate")]
125extern crate accelerate_src;
126
127pub trait ToUsize2 {
128    fn to_usize2(self) -> (usize, usize);
129}
130
131impl ToUsize2 for usize {
132    fn to_usize2(self) -> (usize, usize) {
133        (self, self)
134    }
135}
136
137impl ToUsize2 for (usize, usize) {
138    fn to_usize2(self) -> (usize, usize) {
139        self
140    }
141}
142
143/// Defining a module with forward method using a single argument.
144pub trait Module {
145    fn forward(&self, xs: &Tensor) -> Result<Tensor>;
146}
147
148impl<T: Fn(&Tensor) -> Result<Tensor>> Module for T {
149    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
150        self(xs)
151    }
152}
153
154impl<M: Module> Module for Option<&M> {
155    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
156        match self {
157            None => Ok(xs.clone()),
158            Some(m) => m.forward(xs),
159        }
160    }
161}
162
163/// A single forward method using a single single tensor argument and a flag to
164/// separate the training and evaluation behaviors.
165pub trait ModuleT {
166    fn forward_t(&self, xs: &Tensor, train: bool) -> Result<Tensor>;
167}
168
169impl<M: Module> ModuleT for M {
170    fn forward_t(&self, xs: &Tensor, _train: bool) -> Result<Tensor> {
171        self.forward(xs)
172    }
173}
candle_core/lib.rs

candle_core/
lib.rs
