luma_tensor/ops/construct/
options.rs

1use std::borrow::Cow;
2
3use crate::{
4    Bool, DTypeKind, Device, Float, Int, Layout, Shape, Tensor,
5    dtype::{BoolDType, FloatDType, IntDType},
6};
7
8pub struct TensorCreationOptions<D: Device, K: DTypeKind<D>> {
9    pub device: D,
10    pub dtype: K::DType,
11}
12
13impl<D: Device, K: DTypeKind<D>> From<&D> for TensorCreationOptions<D, K> {
14    fn from(device: &D) -> Self {
15        Self { device: device.clone(), dtype: K::DType::default() }
16    }
17}
18
19impl<D: Device, K: DTypeKind<D>> From<D> for TensorCreationOptions<D, K> {
20    fn from(device: D) -> Self {
21        Self { device, dtype: K::DType::default() }
22    }
23}
24
25impl<D: Device> From<FloatDType> for TensorCreationOptions<D, Float> {
26    fn from(dtype: FloatDType) -> Self {
27        Self { device: D::default(), dtype }
28    }
29}
30
31impl<D: Device> From<IntDType> for TensorCreationOptions<D, Int> {
32    fn from(dtype: IntDType) -> Self {
33        Self { device: D::default(), dtype }
34    }
35}
36
37impl<D: Device> From<BoolDType> for TensorCreationOptions<D, Bool> {
38    fn from(dtype: BoolDType) -> Self {
39        Self { device: D::default(), dtype }
40    }
41}
42
43impl<D: Device, K: DTypeKind<D>> From<()> for TensorCreationOptions<D, K> {
44    fn from(_: ()) -> Self {
45        Self { device: D::default(), dtype: K::DType::default() }
46    }
47}
48
49impl<D: Device, K: DTypeKind<D>> From<(&D, K::DType)> for TensorCreationOptions<D, K> {
50    fn from((device, dtype): (&D, K::DType)) -> Self {
51        Self { device: device.clone(), dtype }
52    }
53}
54
55impl<D: Device, K: DTypeKind<D>> From<(D, K::DType)> for TensorCreationOptions<D, K> {
56    fn from((device, dtype): (D, K::DType)) -> Self {
57        Self { device, dtype }
58    }
59}
60
61pub trait ConstructDTypeKind<D: Device>: DTypeKind<D> {
62    fn zeros_dispatch(shape: &Shape, device: &D, dtype: Self::DType) -> crate::Result<Self::Storage>;
63    fn ones_dispatch(shape: &Shape, device: &D, dtype: Self::DType) -> crate::Result<Self::Storage>;
64    fn full_dispatch(shape: &Shape, value: Self::Scalar, device: &D, dtype: Self::DType) -> crate::Result<Self::Storage>;
65}
66
67impl<D: Device> ConstructDTypeKind<D> for Float {
68    fn zeros_dispatch(shape: &Shape, device: &D, dtype: FloatDType) -> crate::Result<Self::Storage> {
69        D::f_zeros(shape, device, dtype)
70    }
71
72    fn ones_dispatch(shape: &Shape, device: &D, dtype: FloatDType) -> crate::Result<Self::Storage> {
73        D::f_ones(shape, device, dtype)
74    }
75
76    fn full_dispatch(shape: &Shape, value: f64, device: &D, dtype: FloatDType) -> crate::Result<Self::Storage> {
77        D::f_full(shape, value, device, dtype)
78    }
79}
80
81impl<D: Device> ConstructDTypeKind<D> for Int {
82    fn zeros_dispatch(shape: &Shape, device: &D, dtype: IntDType) -> crate::Result<Self::Storage> {
83        D::i_zeros(shape, device, dtype)
84    }
85
86    fn ones_dispatch(shape: &Shape, device: &D, dtype: IntDType) -> crate::Result<Self::Storage> {
87        D::i_ones(shape, device, dtype)
88    }
89
90    fn full_dispatch(shape: &Shape, value: Self::Scalar, device: &D, dtype: IntDType) -> crate::Result<Self::Storage> {
91        D::i_full(shape, value, device, dtype)
92    }
93}
94
95impl<D: Device, K: DTypeKind<D>> Tensor<D, K> {
96    pub fn phantom<S: Into<Shape>>(shape: S, options: impl Into<TensorCreationOptions<D, K>>) -> crate::Result<Self> {
97        let options: TensorCreationOptions<D, K> = options.into();
98        let shape = shape.into();
99        Ok(Self::phantom_storage(shape, options.dtype, options.device))
100    }
101}
102
103impl<D: Device, K: ConstructDTypeKind<D>> Tensor<D, K> {
104    pub fn zeros<S: Into<Shape>>(shape: S, options: impl Into<TensorCreationOptions<D, K>>) -> crate::Result<Self> {
105        let options: TensorCreationOptions<D, K> = options.into();
106        let shape = shape.into();
107        let storage = K::zeros_dispatch(&shape, &options.device, options.dtype)?;
108        Ok(Self::from_storage(storage, shape, K::Meta::default()))
109    }
110
111    pub fn ones<S: Into<Shape>>(shape: S, options: impl Into<TensorCreationOptions<D, K>>) -> crate::Result<Self> {
112        let options: TensorCreationOptions<D, K> = options.into();
113        let shape = shape.into();
114        let storage = K::ones_dispatch(&shape, &options.device, options.dtype)?;
115        Ok(Self::from_storage(storage, shape, K::Meta::default()))
116    }
117
118    pub fn full<S: Into<Shape>>(shape: S, value: K::Scalar, options: impl Into<TensorCreationOptions<D, K>>) -> crate::Result<Self> {
119        let options: TensorCreationOptions<D, K> = options.into();
120        let shape = shape.into();
121        let storage = K::full_dispatch(&shape, value, &options.device, options.dtype)?;
122        Ok(Self::from_storage(storage, shape, K::Meta::default()))
123    }
124
125    pub fn zeros_like(&self) -> crate::Result<Self> {
126        Self::zeros(self.shape().clone(), (self.device(), self.dtype()))
127    }
128
129    pub fn ones_like(&self) -> crate::Result<Self> {
130        Self::ones(self.shape().clone(), (self.device(), self.dtype()))
131    }
132}
133
134// ============================================================================
135//    BytesDTypeKind: kind dispatch for raw-byte I/O
136// ============================================================================
137
138/// Kind-level dispatch for `from_bytes` / `to_bytes`.
139///
140/// Each impl forwards to the corresponding method on the device ops trait
141/// (`FloatOps::f_from_bytes` / `f_to_bytes`, etc.).
142pub trait BytesDTypeKind<D: Device>: DTypeKind<D> {
143    fn from_bytes_dispatch(bytes: &[u8], shape: &Shape, device: &D, dtype: Self::DType) -> crate::Result<Self::Storage>;
144    fn to_bytes_dispatch<'a>(storage: &'a Self::Storage, layout: &Layout) -> crate::Result<Cow<'a, [u8]>>;
145}
146
147impl<D: Device> BytesDTypeKind<D> for Float {
148    fn from_bytes_dispatch(bytes: &[u8], shape: &Shape, device: &D, dtype: FloatDType) -> crate::Result<Self::Storage> {
149        D::f_from_bytes(bytes, shape, device, dtype)
150    }
151
152    fn to_bytes_dispatch<'a>(storage: &'a Self::Storage, layout: &Layout) -> crate::Result<Cow<'a, [u8]>> {
153        D::f_to_bytes(storage, layout)
154    }
155}
156
157impl<D: Device> BytesDTypeKind<D> for Int {
158    fn from_bytes_dispatch(bytes: &[u8], shape: &Shape, device: &D, dtype: IntDType) -> crate::Result<Self::Storage> {
159        D::i_from_bytes(bytes, shape, device, dtype)
160    }
161
162    fn to_bytes_dispatch<'a>(storage: &'a Self::Storage, layout: &Layout) -> crate::Result<Cow<'a, [u8]>> {
163        D::i_to_bytes(storage, layout)
164    }
165}
166
167impl<D: Device> BytesDTypeKind<D> for Bool {
168    fn from_bytes_dispatch(bytes: &[u8], shape: &Shape, device: &D, _dtype: BoolDType) -> crate::Result<Self::Storage> {
169        D::b_from_bytes(bytes, shape, device, BoolDType::Bool)
170    }
171
172    fn to_bytes_dispatch<'a>(storage: &'a Self::Storage, layout: &Layout) -> crate::Result<Cow<'a, [u8]>> {
173        D::b_to_bytes(storage, layout)
174    }
175}
176
177impl<D: Device, K: BytesDTypeKind<D>> Tensor<D, K> {
178    /// Create a tensor from raw little-endian bytes.
179    ///
180    /// The byte slice length must equal `shape.element_count() * dtype.size_in_bytes()`.
181    pub fn from_bytes<'a>(
182        bytes: impl Into<Cow<'a, [u8]>>,
183        shape: impl Into<Shape>,
184        options: impl Into<TensorCreationOptions<D, K>>,
185    ) -> crate::Result<Self> {
186        let options: TensorCreationOptions<D, K> = options.into();
187        let shape = shape.into();
188        let storage = K::from_bytes_dispatch(&bytes.into(), &shape, &options.device, options.dtype)?;
189        Ok(Self::from_storage(storage, shape, K::Meta::default()))
190    }
191
192    /// Read raw little-endian bytes in logical (layout) order.
193    ///
194    /// Always returns owned bytes (the internal `RwLock` prevents zero-copy
195    /// borrowing at this level; the device trait internally uses `Cow` for
196    /// contiguous-optimised paths).
197    pub fn to_bytes(&self) -> crate::Result<Vec<u8>> {
198        let guard = self.storage_read()?;
199        Ok(K::to_bytes_dispatch(&*guard, self.layout())?.into_owned())
200    }
201}
luma_tensor/ops/construct/options.rs

luma_tensor/ops/construct/
options.rs
