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}