tenferro_tensor/tensor/constructors/
deterministic.rs

1use std::sync::Arc;
2
3use tenferro_algebra::Scalar;
4use tenferro_device::{Error, LogicalMemorySpace, Result};
5
6use super::super::{Tensor, TensorParts};
7use crate::layout::validate_layout_against_len;
8use crate::{DataBuffer, MemoryOrder};
9
10impl<T: Scalar> Tensor<T> {
11    pub(crate) fn finish_allocation(
12        tensor: Self,
13        memory_space: LogicalMemorySpace,
14    ) -> Result<Self> {
15        if memory_space == LogicalMemorySpace::MainMemory {
16            Ok(tensor)
17        } else {
18            tensor.to_memory_space_async(memory_space)
19        }
20    }
21
22    pub(crate) fn main_memory_contiguous(data: Vec<T>, dims: &[usize], order: MemoryOrder) -> Self {
23        Self::from_owned_contiguous_data(
24            data,
25            Arc::from(dims),
26            order,
27            LogicalMemorySpace::MainMemory,
28            None,
29            false,
30        )
31    }
32
33    /// Layout policy for the `*_like` constructors.
34    ///
35    /// Row-major is preserved only when the source is row-major contiguous
36    /// and not column-major contiguous. Ambiguous or non-contiguous sources
37    /// fall back to column-major so the result layout remains deterministic.
38    pub(super) fn like_order(reference: &Self) -> MemoryOrder {
39        if reference.is_row_major_contiguous() && !reference.is_col_major_contiguous() {
40            MemoryOrder::RowMajor
41        } else {
42            MemoryOrder::ColumnMajor
43        }
44    }
45
46    fn required_storage_len(dims: &[usize], strides: &[isize], offset: isize) -> Result<usize> {
47        if dims.len() != strides.len() {
48            return Err(Error::InvalidArgument(format!(
49                "empty_strided rank mismatch: dims={} strides={}",
50                dims.len(),
51                strides.len()
52            )));
53        }
54
55        if dims.iter().product::<usize>() == 0 {
56            return Ok(0);
57        }
58
59        let mut min_pos = offset;
60        let mut max_pos = offset;
61        for (axis, (&dim, &stride)) in dims.iter().zip(strides).enumerate() {
62            let extent = isize::try_from(dim - 1)
63                .ok()
64                .and_then(|d| d.checked_mul(stride))
65                .ok_or_else(|| {
66                    Error::StrideError(format!(
67                        "empty_strided extent overflow for dimension {axis} (size={dim}, stride={stride})"
68                    ))
69                })?;
70            if extent >= 0 {
71                max_pos = max_pos.checked_add(extent).ok_or_else(|| {
72                    Error::StrideError(format!(
73                        "empty_strided maximum offset overflow for dimension {axis}"
74                    ))
75                })?;
76            } else {
77                min_pos = min_pos.checked_add(extent).ok_or_else(|| {
78                    Error::StrideError(format!(
79                        "empty_strided minimum offset overflow for dimension {axis}"
80                    ))
81                })?;
82            }
83        }
84
85        if min_pos < 0 {
86            return Err(Error::StrideError(format!(
87                "empty_strided accesses negative buffer positions {}..={}",
88                min_pos, max_pos
89            )));
90        }
91
92        let max_pos = usize::try_from(max_pos).map_err(|_| {
93            Error::StrideError(format!(
94                "empty_strided maximum position {max_pos} exceeds usize range"
95            ))
96        })?;
97        max_pos
98            .checked_add(1)
99            .ok_or_else(|| Error::StrideError("empty_strided storage length overflow".into()))
100    }
101
102    /// Create a tensor filled with zeros.
103    ///
104    /// Allocates directly on the target device without an intermediate CPU
105    /// buffer. For GPU targets (with the `cuda` feature) this avoids the
106    /// CPU-allocate-then-transfer overhead.
107    ///
108    /// # Examples
109    ///
110    /// ```ignore
111    /// use tenferro_tensor::{MemoryOrder, Tensor};
112    /// use tenferro_device::LogicalMemorySpace;
113    ///
114    /// let a = Tensor::<f64>::zeros(
115    ///     &[3, 4],
116    ///     LogicalMemorySpace::MainMemory,
117    ///     MemoryOrder::ColumnMajor,
118    /// ).unwrap();
119    /// ```
120    pub fn zeros(
121        dims: &[usize],
122        memory_space: LogicalMemorySpace,
123        order: MemoryOrder,
124    ) -> Result<Self> {
125        let n_elements: usize = dims.iter().product();
126        let buffer = DataBuffer::zeros_on_device(n_elements, memory_space)?;
127        let strides = crate::layout::compute_contiguous_strides(dims, order);
128        Ok(Self::from_parts(TensorParts {
129            buffer,
130            dims: Arc::from(dims),
131            strides: Arc::from(strides),
132            offset: 0,
133            logical_memory_space: memory_space,
134            preferred_compute_device: None,
135            event: None,
136            conjugated: false,
137            fw_grad: None,
138        }))
139    }
140
141    /// Create a tensor with allocated storage.
142    ///
143    /// The current safe implementation initializes the backing storage to
144    /// zero, which keeps the constructor deterministic while preserving the
145    /// requested layout and device placement. For GPU targets this allocates
146    /// uninitialised device memory directly.
147    ///
148    /// The `*_like` family preserves a row-major layout only when the source
149    /// tensor is row-major contiguous and not column-major contiguous.
150    /// Ambiguous or non-contiguous inputs fall back to column-major.
151    ///
152    /// # Examples
153    ///
154    /// ```ignore
155    /// use tenferro_tensor::{MemoryOrder, Tensor};
156    /// use tenferro_device::LogicalMemorySpace;
157    ///
158    /// let a = Tensor::<f64>::empty(
159    ///     &[3, 4],
160    ///     LogicalMemorySpace::MainMemory,
161    ///     MemoryOrder::ColumnMajor,
162    /// ).unwrap();
163    /// ```
164    pub fn empty(
165        dims: &[usize],
166        memory_space: LogicalMemorySpace,
167        order: MemoryOrder,
168    ) -> Result<Self> {
169        let n_elements: usize = dims.iter().product();
170        let buffer = DataBuffer::allocate_uninit_on_device(n_elements, memory_space)?;
171        let strides = crate::layout::compute_contiguous_strides(dims, order);
172        Ok(Self::from_parts(TensorParts {
173            buffer,
174            dims: Arc::from(dims),
175            strides: Arc::from(strides),
176            offset: 0,
177            logical_memory_space: memory_space,
178            preferred_compute_device: None,
179            event: None,
180            conjugated: false,
181            fw_grad: None,
182        }))
183    }
184
185    /// Create a tensor filled with ones.
186    ///
187    /// # Examples
188    ///
189    /// ```ignore
190    /// use tenferro_tensor::{MemoryOrder, Tensor};
191    /// use tenferro_device::LogicalMemorySpace;
192    ///
193    /// let a = Tensor::<f64>::ones(
194    ///     &[2, 3],
195    ///     LogicalMemorySpace::MainMemory,
196    ///     MemoryOrder::ColumnMajor,
197    /// ).unwrap();
198    /// ```
199    pub fn ones(
200        dims: &[usize],
201        memory_space: LogicalMemorySpace,
202        order: MemoryOrder,
203    ) -> Result<Self> {
204        let n_elements: usize = dims.iter().product();
205        Self::finish_allocation(
206            Self::main_memory_contiguous(vec![T::one(); n_elements], dims, order),
207            memory_space,
208        )
209    }
210
211    /// Create a tensor with explicit strides.
212    ///
213    /// # Errors
214    ///
215    /// Returns an error if the layout would access storage outside the
216    /// allocated buffer.
217    ///
218    /// # Examples
219    ///
220    /// ```ignore
221    /// use tenferro_tensor::Tensor;
222    /// use tenferro_device::LogicalMemorySpace;
223    ///
224    /// let t = Tensor::<f64>::empty_strided(&[2, 2], &[1, 2], 0, LogicalMemorySpace::MainMemory).unwrap();
225    /// assert_eq!(t.strides(), &[1, 2]);
226    /// ```
227    pub fn empty_strided(
228        dims: &[usize],
229        strides: &[isize],
230        offset: isize,
231        memory_space: LogicalMemorySpace,
232    ) -> Result<Self> {
233        let storage_len = Self::required_storage_len(dims, strides, offset)?;
234        // Allocate zeros (not uninit) because the storage may contain gaps
235        // between strided elements that should not hold garbage.
236        let buffer = DataBuffer::zeros_on_device(storage_len, memory_space)?;
237        validate_layout_against_len(dims, strides, offset, buffer.len())?;
238        Ok(Self::from_parts(TensorParts {
239            buffer,
240            dims: Arc::from(dims),
241            strides: Arc::from(strides),
242            offset,
243            logical_memory_space: memory_space,
244            preferred_compute_device: None,
245            event: None,
246            conjugated: false,
247            fw_grad: None,
248        }))
249    }
250
251    /// Create a tensor filled with `value`.
252    ///
253    /// # Examples
254    ///
255    /// ```ignore
256    /// use tenferro_tensor::{MemoryOrder, Tensor};
257    /// use tenferro_device::LogicalMemorySpace;
258    ///
259    /// let a = Tensor::<f64>::full(
260    ///     &[2, 3],
261    ///     7.5,
262    ///     LogicalMemorySpace::MainMemory,
263    ///     MemoryOrder::ColumnMajor,
264    /// ).unwrap();
265    /// ```
266    pub fn full(
267        dims: &[usize],
268        value: T,
269        memory_space: LogicalMemorySpace,
270        order: MemoryOrder,
271    ) -> Result<Self> {
272        let n_elements: usize = dims.iter().product();
273        Self::finish_allocation(
274            Self::main_memory_contiguous(vec![value; n_elements], dims, order),
275            memory_space,
276        )
277    }
278
279    /// Create a tensor from a data slice.
280    ///
281    /// `order` describes how to interpret `data` at the import boundary. View
282    /// operations continue to use tenferro's internal column-major semantics.
283    ///
284    /// # Errors
285    ///
286    /// Returns an error if `data.len()` does not match the product of `dims`.
287    ///
288    /// # Examples
289    ///
290    /// ```ignore
291    /// use tenferro_tensor::{MemoryOrder, Tensor};
292    ///
293    /// let data = [1.0, 2.0, 3.0, 4.0];
294    /// let t = Tensor::<f64>::from_slice(&data, &[2, 2], MemoryOrder::ColumnMajor).unwrap();
295    /// ```
296    pub fn from_slice(data: &[T], dims: &[usize], order: MemoryOrder) -> Result<Self> {
297        let n_elements: usize = dims.iter().product();
298        if data.len() != n_elements {
299            return Err(Error::InvalidArgument(format!(
300                "data length {} doesn't match dims product {}",
301                data.len(),
302                n_elements
303            )));
304        }
305        Ok(Self::main_memory_contiguous(data.to_vec(), dims, order))
306    }
307
308    /// Create a tensor from a row-major data slice.
309    ///
310    /// This is a convenience wrapper around
311    /// [`from_slice`](Self::from_slice) with `MemoryOrder::RowMajor`.
312    /// It lets NumPy / C users pass data in their natural order while
313    /// tenferro internally stores it in column-major layout.
314    ///
315    /// # Errors
316    ///
317    /// Returns an error if `data.len()` does not match the product of `dims`.
318    ///
319    /// # Examples
320    ///
321    /// ```
322    /// use tenferro_tensor::{MemoryOrder, Tensor};
323    ///
324    /// // Row-major: data is laid out row by row.
325    /// // [[1, 2],
326    /// //  [3, 4]]
327    /// let t = Tensor::<f64>::from_row_major_slice(
328    ///     &[1.0, 2.0, 3.0, 4.0],
329    ///     &[2, 2],
330    /// ).unwrap();
331    /// assert_eq!(t.dims(), &[2, 2]);
332    /// ```
333    pub fn from_row_major_slice(data: &[T], dims: &[usize]) -> Result<Self> {
334        Self::from_slice(data, dims, MemoryOrder::RowMajor)
335    }
336
337    /// Create a tensor from an owned `Vec<T>` with explicit layout.
338    ///
339    /// # Errors
340    ///
341    /// Returns an error if the layout is inconsistent with the data length.
342    ///
343    /// # Examples
344    ///
345    /// ```ignore
346    /// use tenferro_tensor::Tensor;
347    ///
348    /// let t = Tensor::<f64>::from_vec(vec![1.0, 2.0, 3.0, 4.0], &[2, 2], &[1, 2], 0).unwrap();
349    /// ```
350    pub fn from_vec(
351        data: Vec<T>,
352        dims: &[usize],
353        strides: &[isize],
354        offset: isize,
355    ) -> Result<Self> {
356        validate_layout_against_len(dims, strides, offset, data.len())?;
357        Ok(Self::from_parts(TensorParts {
358            buffer: DataBuffer::from_vec(data),
359            dims: Arc::from(dims),
360            strides: Arc::from(strides),
361            offset,
362            logical_memory_space: LogicalMemorySpace::MainMemory,
363            preferred_compute_device: None,
364            event: None,
365            conjugated: false,
366            fw_grad: None,
367        }))
368    }
369
370    /// Create a tensor from externally-owned CPU-accessible memory.
371    ///
372    /// # Safety
373    ///
374    /// - `ptr` must remain valid for at least `len` elements until `release` is called.
375    /// - The layout described by `dims`, `strides`, and `offset` must stay in bounds.
376    ///
377    /// # Examples
378    ///
379    /// ```ignore
380    /// use tenferro_tensor::Tensor;
381    ///
382    /// let data = vec![1.0, 2.0, 3.0, 4.0];
383    /// let ptr = data.as_ptr();
384    /// let tensor = unsafe {
385    ///     Tensor::from_external_parts(ptr, data.len(), &[2, 2], &[1, 2], 0, move || drop(data))
386    /// }.unwrap();
387    /// assert_eq!(tensor.dims(), &[2, 2]);
388    /// ```
389    pub unsafe fn from_external_parts(
390        ptr: *const T,
391        len: usize,
392        dims: &[usize],
393        strides: &[isize],
394        offset: isize,
395        release: impl FnOnce() + Send + 'static,
396    ) -> Result<Self> {
397        validate_layout_against_len(dims, strides, offset, len)?;
398        Ok(Self::from_parts(TensorParts {
399            buffer: DataBuffer::from_external(ptr, len, release),
400            dims: Arc::from(dims),
401            strides: Arc::from(strides),
402            offset,
403            logical_memory_space: LogicalMemorySpace::MainMemory,
404            preferred_compute_device: None,
405            event: None,
406            conjugated: false,
407            fw_grad: None,
408        }))
409    }
410
411    /// Try to extract the underlying data as `Vec<T>`.
412    ///
413    /// # Examples
414    ///
415    /// ```ignore
416    /// let t = Tensor::<f64>::from_slice(&[1.0, 2.0], &[2], MemoryOrder::ColumnMajor).unwrap();
417    /// let _data = t.try_into_data_vec();
418    /// ```
419    pub fn try_into_data_vec(self) -> Option<Vec<T>> {
420        self.buffer.try_into_vec()
421    }
422
423    /// Create a tensor with the same shape and layout convention as another tensor.
424    ///
425    /// # Examples
426    ///
427    /// ```ignore
428    /// use tenferro_tensor::{MemoryOrder, Tensor};
429    /// use tenferro_device::LogicalMemorySpace;
430    ///
431    /// let base = Tensor::<f64>::zeros(
432    ///     &[2, 3],
433    ///     LogicalMemorySpace::MainMemory,
434    ///     MemoryOrder::ColumnMajor,
435    /// ).unwrap();
436    /// let like = base.empty_like().unwrap();
437    /// assert_eq!(like.dims(), base.dims());
438    /// ```
439    pub fn empty_like(&self) -> Result<Self> {
440        Self::empty(
441            self.dims(),
442            self.logical_memory_space(),
443            Self::like_order(self),
444        )
445    }
446
447    /// Create a zero-filled tensor with the same shape and layout convention as another tensor.
448    ///
449    /// The `*_like` family preserves a row-major layout only when the source
450    /// tensor is row-major contiguous and not column-major contiguous.
451    /// Ambiguous or non-contiguous inputs fall back to column-major.
452    ///
453    /// # Examples
454    ///
455    /// ```ignore
456    /// use tenferro_tensor::{MemoryOrder, Tensor};
457    /// use tenferro_device::LogicalMemorySpace;
458    ///
459    /// let base = Tensor::<f64>::zeros(
460    ///     &[2, 3],
461    ///     LogicalMemorySpace::MainMemory,
462    ///     MemoryOrder::ColumnMajor,
463    /// ).unwrap();
464    /// let like = base.zeros_like().unwrap();
465    /// assert_eq!(like.dims(), base.dims());
466    /// ```
467    pub fn zeros_like(&self) -> Result<Self> {
468        Self::zeros(
469            self.dims(),
470            self.logical_memory_space(),
471            Self::like_order(self),
472        )
473    }
474
475    /// Create a one-filled tensor with the same shape and layout convention as another tensor.
476    ///
477    /// The `*_like` family preserves a row-major layout only when the source
478    /// tensor is row-major contiguous and not column-major contiguous.
479    /// Ambiguous or non-contiguous inputs fall back to column-major.
480    ///
481    /// # Examples
482    ///
483    /// ```ignore
484    /// use tenferro_tensor::{MemoryOrder, Tensor};
485    /// use tenferro_device::LogicalMemorySpace;
486    ///
487    /// let base = Tensor::<f64>::zeros(
488    ///     &[2, 3],
489    ///     LogicalMemorySpace::MainMemory,
490    ///     MemoryOrder::ColumnMajor,
491    /// ).unwrap();
492    /// let like = base.ones_like().unwrap();
493    /// assert_eq!(like.dims(), base.dims());
494    /// ```
495    pub fn ones_like(&self) -> Result<Self> {
496        Self::ones(
497            self.dims(),
498            self.logical_memory_space(),
499            Self::like_order(self),
500        )
501    }
502
503    /// Create a tensor filled with `value` and matching the shape/layout convention of another tensor.
504    ///
505    /// The `*_like` family preserves a row-major layout only when the source
506    /// tensor is row-major contiguous and not column-major contiguous.
507    /// Ambiguous or non-contiguous inputs fall back to column-major.
508    ///
509    /// # Examples
510    ///
511    /// ```ignore
512    /// use tenferro_tensor::{MemoryOrder, Tensor};
513    /// use tenferro_device::LogicalMemorySpace;
514    ///
515    /// let base = Tensor::<f64>::zeros(
516    ///     &[2, 3],
517    ///     LogicalMemorySpace::MainMemory,
518    ///     MemoryOrder::ColumnMajor,
519    /// ).unwrap();
520    /// let like = base.full_like(3.25).unwrap();
521    /// assert_eq!(like.dims(), base.dims());
522    /// ```
523    pub fn full_like(&self, value: T) -> Result<Self> {
524        Self::full(
525            self.dims(),
526            value,
527            self.logical_memory_space(),
528            Self::like_order(self),
529        )
530    }
531}