ariadnetor_tensor/
lib.rs

1//! Tensor storage library
2//!
3//! Provides backend-agnostic data structures for tensor storage and
4//! the user-facing [`Tensor<St, L>`] type that joins them. The tensor
5//! carries no backend; operations take it explicitly at the call site.
6
7#![deny(missing_docs)]
8
9mod block_sparse;
10mod capability;
11mod dense;
12mod error;
13mod layout;
14mod ops;
15mod reorder;
16mod sector;
17mod storage;
18mod tensor;
19mod tensor_data;
20
21// Shared BlockSparse test fixtures: `QNIndex` leg builders and densify /
22// template helpers. Reachable by this crate's own unit tests under
23// `cfg(test)`, and by other crates' tests via the `test-fixtures`
24// dev-dependency feature. See the module docs for why a separate crate cannot
25// serve the in-lib unit tests.
26#[cfg(any(test, feature = "test-fixtures"))]
27pub mod test_fixtures;
28
29// Re-export from ariadnetor-core
30pub use ariadnetor_core::{
31    Complex, ComputeBackend, ContractionError, ContractionPlan, EinsumExpr, MemoryOrder, Scalar,
32    compute_permutation,
33};
34
35pub use block_sparse::{
36    BlockCoord, BlockMeta, BlockSparseLayout, BlockSparseStorage, BlockSparseTensorData, Direction,
37    QNIndex,
38};
39pub use capability::{Host, OpsFor};
40pub use dense::{DenseLayout, DenseStorage, DenseTensorData};
41pub use error::TensorError;
42pub use layout::{StorageFor, TensorLayout};
43pub use ops::{add_all, linear_combine};
44pub use reorder::{flat_index, normalize_to_data, reorder_data};
45pub use sector::{Sector, U1Sector, Z2Sector};
46pub use storage::Storage;
47pub use tensor::{BlockSparseTensor, DenseTensor, Tensor};
48pub use tensor_data::TensorData;
49
50// Re-export the native backend so users importing ariadnetor_tensor for
51// `DenseTensor` / `BlockSparseTensor` constructors get its default
52// backend without a separate crate import.
53pub use ariadnetor_native::NativeBackend;
54
55/// Extension trait for backend-aware tensor construction.
56///
57/// Provides tensor constructors on any `ComputeBackend`. The constructed
58/// tensors have `order()` matching the backend's `preferred_order()`, so
59/// downstream linalg operations driven by that backend find the storage
60/// already in the layout they expect. Most constructors return the
61/// Mid-layer `DenseTensorData<T>` for kernel-output paths; `dense`
62/// returns the wrapped `DenseTensor<T>` for the input-fabrication case.
63pub trait ComputeBackendTensorExt: ComputeBackend {
64    /// Construct a `DenseTensorData` from data in this backend's
65    /// preferred memory order.
66    ///
67    /// The caller must arrange `data` in this backend's
68    /// `preferred_order()`. The resulting tensor data has
69    /// `order() == self.preferred_order()`.
70    fn make_tensor<T: Clone>(&self, data: Vec<T>, shape: Vec<usize>) -> DenseTensorData<T> {
71        DenseTensorData::from_raw_parts(data, shape, self.preferred_order())
72    }
73
74    /// Construct a `DenseTensor` from data in this backend's preferred
75    /// memory order.
76    ///
77    /// The caller must arrange `data` in this backend's
78    /// `preferred_order()`. The resulting tensor has
79    /// `order() == self.preferred_order()`.
80    ///
81    /// One-call entry for fabricating an input tensor from flat parts:
82    /// fuses [`make_tensor`](Self::make_tensor) — which yields the
83    /// Mid-layer `DenseTensorData`, kept for kernel-output paths — with
84    /// the `DenseTensor` wrap, so a caller that just wants a tensor need
85    /// not reach across the Data layer. The backend stays explicit at
86    /// the call site, so this is not a host-hardcoded constructor.
87    fn dense<T: Clone>(&self, data: Vec<T>, shape: Vec<usize>) -> DenseTensor<T> {
88        DenseTensor::from_data(self.make_tensor(data, shape))
89    }
90
91    /// Create a zero-filled tensor whose `order()` matches this backend.
92    fn zeros<T: Clone + num_traits::Zero>(&self, shape: Vec<usize>) -> DenseTensorData<T> {
93        DenseTensorData::zeros_in_order(shape, self.preferred_order())
94    }
95
96    /// Create a ones-filled tensor whose `order()` matches this backend.
97    fn ones<T: Clone + num_traits::Zero + num_traits::One>(
98        &self,
99        shape: Vec<usize>,
100    ) -> DenseTensorData<T> {
101        DenseTensorData::ones_in_order(shape, self.preferred_order())
102    }
103
104    /// Create a tensor filled with `value` whose `order()` matches this backend.
105    fn filled<T: Clone>(&self, shape: Vec<usize>, value: T) -> DenseTensorData<T> {
106        DenseTensorData::filled_in_order(shape, value, self.preferred_order())
107    }
108
109    /// Create an identity matrix whose `order()` matches this backend.
110    ///
111    /// The identity matrix is symmetric, so its flat data layout is the same
112    /// regardless of memory order; only the `order()` field differs.
113    fn eye<T: Clone + num_traits::Zero + num_traits::One>(&self, n: usize) -> DenseTensorData<T> {
114        DenseTensorData::eye_in_order(n, self.preferred_order())
115    }
116}
117
118impl<B: ComputeBackend> ComputeBackendTensorExt for B {}
ariadnetor_tensor/lib.rs

ariadnetor_tensor/
lib.rs
