forge/
lib.rs

1//! Forge — a WebGPU-native machine learning framework in Rust, scoped to GPT-2.
2//!
3//! Production execution targets WebGPU via `wgpu`; the CPU backend is a
4//! mathematically identical reference used for testing and verification.
5//!
6//! # Optional features
7//!
8//! One, off by default, so `cargo add forge-ml` gets the runtime alone.
9//!
10//! - **`train`** — [`autograd`] and [`optim`], the backward kernels, and
11//!   `Gpt2::loss`/`loss_grads`. Off by default because Forge is an inference
12//!   runtime that also happens to train; `cargo add forge-ml` should not
13//!   compile a tape you never record. Construction and serialization
14//!   (`Gpt2::init_random`, `params`, `save_safetensors`) stay core — they are
15//!   not training, and the inference tests use them.
16
17// Nightly-only, and only ever set by docs.rs (see `[package.metadata.docs.rs]`),
18// so this is inert on stable. It is what puts the "Available on crate feature"
19// badges on the items below.
20#![cfg_attr(docsrs, feature(doc_cfg))]
21
22// Behind `train`. The module's own `//!` docs carry the description; an outer
23// `///` here would move intra-doc link resolution into this file's scope and
24// break `[`Tape`]`.
25#[cfg(feature = "train")]
26#[cfg_attr(docsrs, doc(cfg(feature = "train")))]
27pub mod autograd;
28pub mod backend;
29pub mod device;
30pub mod dtype;
31pub mod error;
32pub mod models;
33pub mod nn;
34pub mod ops;
35// Behind `train`; see the note on `autograd` above.
36#[cfg(feature = "train")]
37#[cfg_attr(docsrs, doc(cfg(feature = "train")))]
38pub mod optim;
39pub mod serialization;
40pub mod shape;
41pub mod tensor;
42pub mod tokenizer;
43#[cfg(target_arch = "wasm32")]
44pub mod wasm;
45
46pub use device::Device;
47pub use dtype::DType;
48pub use error::{ForgeError, Result};
49pub use models::gpt2::{
50    AttnStep, Gpt2, Gpt2Config, KvCache, LayerDetail, Sampler, Sampling, StepTrace, Surprisal,
51    top_probs,
52};
53pub use shape::Shape;
54pub use tensor::Tensor;
55pub use tokenizer::{AnyTokenizer, CharTokenizer, Gpt2Tokenizer, Tokenizer};
forge/lib.rs

forge/
lib.rs
