oarust 0.8.6

GPU-first semantic computing engine written in Rust
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55

//! Machine-learning operations, differentiation, layers, and optimizers.

mod actor_critic;
pub mod advantage;
mod autograd;
pub mod byte;
mod checkpoint;
mod collector;
pub mod environment;
pub mod evaluation;
pub mod flow;
pub mod loss;
pub mod matrix;
pub mod mcp_training;
pub mod metric;
mod model_file;
mod module;
pub mod nn;
pub mod optim;
mod parameter;
pub mod policy;
mod random;
pub mod replay;
pub mod rollout;
pub mod tokenizer;
pub mod training;
mod validation;
mod weights;

pub use actor_critic::{
	ActorCritic, ActorCriticOutput, CategoricalActorCritic, CategoricalActorCriticConfig,
};
pub use autograd::GradientTape;
pub use checkpoint::{
	CheckpointManager, CheckpointManagerConfig, load_checkpoint, save_checkpoint,
};
pub(crate) use checkpoint::{
	DenseArtifactTensor, ModuleArtifact, ModuleArtifactMetadata, save_dense_artifact,
	save_module_artifact,
};
pub use collector::{RolloutCollector, RolloutCollectorConfig, RolloutCollectorMetrics};
pub use environment::{
	EnvironmentSpace, EnvironmentSpaceKind, EnvironmentSpec, EnvironmentTransition,
};
pub use matrix::UpsampleMode;
pub use mcp_training::{McpTraining, McpTrainingConfig};
pub use module::{Module, ModuleRegistry, NamedBuffer, NamedParameter, NamedStateU32, ScopedEval};
pub use optim::{Adam, AdamW, CheckpointOptimizer, Muon, NoOpOptimizer, Optimizer, Sgd};
pub use parameter::Parameter;
pub use policy::{ContinuousPolicyResult, PolicyResult};
pub use replay::{ReplayBatch, ReplayBuffer, ReplayConfig, ReplayTransition};
pub use rollout::{RolloutBatch, RolloutBuffer, RolloutConfig, RolloutTransition};
pub use tokenizer::{BpeMerge, BpeTokenizer};
pub use training::*;
pub(crate) use weights::SafeTensorsSource;