ruda_optim/
lib.rs

1#![cfg_attr(not(feature = "std"), no_std)]
2#![warn(missing_docs)]
3#![cfg_attr(docsrs, feature(doc_cfg))]
4#![recursion_limit = "256"]
5
6//! Ruda optimizer updates, gradient state, clipping and learning-rate schedules.
7
8#[macro_use]
9extern crate derive_new;
10
11extern crate alloc;
12
13/// Optimizer module.
14mod optim;
15pub use optim::*;
16
17/// Opt-in single-kernel AdamW/AMSGrad and explicit numerical reference.
18#[cfg(feature = "fused-adamw")]
19pub mod fused_adamw;
20
21/// Gradient clipping module.
22pub mod grad_clipping;
23
24/// Learning rate scheduler module.
25#[cfg(feature = "std")]
26pub mod lr_scheduler;
27
28/// Combined model, optimizer, scheduler and accumulated-gradient records.
29#[cfg(feature = "std")]
30pub mod training;
31
32/// Explicit replicated data-parallel training through ruCCL rank communicators.
33#[cfg(feature = "collective")]
34pub mod data_parallel;
35
36/// Type alias for the learning rate.
37///
38/// LearningRate also implements [learning rate scheduler](crate::lr_scheduler::LrScheduler) so it
39/// can be used for constant learning rate.
40pub type LearningRate = f64; // We could potentially change the type.
41
42/// Backend for test cases
43#[cfg(all(
44    test,
45    not(feature = "test-tch"),
46    not(feature = "test-wgpu"),
47    not(feature = "test-cuda"),
48    not(feature = "test-rocm")
49))]
50pub type TestBackend = ruda_tensor_host::Host;
51
52#[cfg(all(test, feature = "test-tch"))]
53/// Backend for test cases
54pub type TestBackend = ruda_tensor_tch::LibTorch<f32>;
55
56#[cfg(all(test, feature = "test-wgpu"))]
57/// Backend for test cases
58pub type TestBackend = ruda_tensor_wgpu::Wgpu;
59
60#[cfg(all(test, feature = "test-cuda"))]
61/// Backend for test cases
62pub type TestBackend = ruda_tensor_device::cuda::Cuda;
63
64#[cfg(all(test, feature = "test-rocm"))]
65/// Backend for test cases
66pub type TestBackend = ruda_tensor_rocm::Rocm;
67
68/// Backend for autodiff test cases
69#[cfg(test)]
70pub type TestAutodiffBackend = ruda_autodiff::Autodiff<TestBackend>;
71
72#[cfg(all(test, feature = "test-memory-checks"))]
73mod tests {
74    ruda_fusion::memory_checks!();
75}
ruda_optim/lib.rs

ruda_optim/
lib.rs
