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