Skip to main content

ruda_tensor/api/
mod.rs

1pub(crate) mod stats;
2
3pub(crate) mod check;
4
5mod autodiff;
6mod base;
7mod bool;
8mod cartesian_grid;
9#[path = "cast.rs"]
10mod tensor_cast;
11mod float;
12mod fmod;
13mod int;
14mod numeric;
15mod moe;
16mod exchange;
17mod moe_exchange;
18pub use moe_exchange::{NativeMoeDispatched,dispatch_moe,combine_moe,received_moe_experts};
19mod options;
20mod orderable;
21mod pad;
22mod spatial_pool;
23mod spatial_interpolate;
24pub use pad::IntoPadding;
25mod take;
26mod transaction;
27pub mod sparse;
28
29mod trunc;
30
31pub use autodiff::*;
32pub use base::*;
33pub use cartesian_grid::cartesian_grid;
34pub use tensor_cast::*;
35pub use crate::element::cast;
36pub use float::{DEFAULT_ATOL, DEFAULT_RTOL};
37pub use numeric::*;
38pub use moe::{selected_router_weights,selected_router_weights_backward};
39pub use options::*;
40pub use transaction::*;
41
42pub use crate::tensor::IndexingUpdateOp;
43pub use crate::collective::{TensorCollective,ReplicatedTensorCollective,BroadcastTensorCollective,IntegerTensorCollective,VariableTensorCollective,VariableTensorExchange};
44pub use crate::frozen_awq::{FrozenAwqOps, FrozenAwqError};
45pub use crate::frozen_nf4::{FrozenNf4Ops,FrozenNf4Error,Nf4ProjectionOptions};
46pub use crate::grouped_nf4::{Nf4GroupedOptions,Nf4ExpertPayload,FrozenNf4GroupedOps,FrozenNf4SwiGluOps};
47pub use crate::packed_experts::{AwqExpertOptions,AwqExpertPayload,PackedExpertPayload,PackedExpertAutodiffError,FrozenPackedExpertOps};
48pub use crate::expert_projection::{ExpertProjectionOptions,ExpertProjectionSelection,ExpertProjectionBackward,ExpertProjectionOps,
49    NativeSwiGluSelection,NativeSwiGluBackward,NativeSwiGluOps};
50pub use crate::moe::{MoeOps,MoeOptions,MoeSelectionOptions,MoeRouterScoring,MoeRouterWeightOptions,
51    MoeExpertStrategy,MoeCombineGradientStrategy,MoeBackward,MoeGradientSelection,MoeBackwardSelected,MoeAutodiffError};
52pub use crate::moe_exchange::{MoeDispatched,MoeDispatchOps,MoeCombineSelection,MoeCombineBackward,
53    MoeReceivedOptions,MoeReceivedSelection,MoeReceivedBackward,MoeReceivedOps};
54
55// Re-exported types
56pub use crate::{
57    BoolDType, BoolStore, DType, DataError, FloatDType, IntDType, TensorData, TensorMetadata,
58    TensorPrimitive, Tolerance,
59    distribution::*,
60    element::*,
61    indexing::*,
62    ops::TransactionPrimitive,
63    shape::*,
64    slice::*,
65    tensor::{Bool, Float, Int, TensorKind},
66};
67
68/// The activation module.
69pub mod activation;
70
71/// The backend module.
72pub mod backend {
73    pub use crate::backend::*;
74}
75
76/// The container module.
77pub mod container {
78    pub use crate::tensor::TensorContainer;
79}
80
81/// The grid module.
82pub mod grid;
83
84/// The linalg module.
85pub mod linalg;
86
87/// The loss module.
88pub mod loss;
89
90/// The neural network module.
91pub mod module;
92
93/// The signal processing module.
94pub mod signal;
95
96/// Operations on tensors module.
97pub mod ops {
98    pub use crate::backend::ops::*;
99    pub use crate::tensor::{
100        BoolElem, BoolTensor, Device, FloatElem, FloatTensor, IntElem, IntTensor, QuantizedTensor,
101    };
102}
103
104/// Tensor quantization module.
105pub mod quantization;
106
107#[cfg(feature = "api-std")]
108pub use report::*;
109
110#[cfg(feature = "api-std")]
111mod report;
112
113pub use ops::Device; // Re-export device so that it's available from `ruda_tensor::api::Device`.
114
115pub(crate) use check::macros::check;
116
117pub use crate::{
118    AllocationProperty, Bytes, DeviceSettings, StreamId, bf16, f16, get_device_settings, read_sync,
119    set_default_dtypes, try_read_sync,
120};