candela/tensor/tensor_interface.rs
1#![allow(private_bounds)]
2use crate::OpError;
3use crate::tensor::backend::{Backend, ComputeFor, DefaultBackend};
4use crate::tensor::graph::{NodeKind, TensorGraphEdge};
5use crate::tensor::iter::{InformedIter, Iter, StepInfo};
6use crate::tensor::mem_formats::layout::Layout;
7use crate::tensor::promise::TensorPromise;
8use crate::tensor::skeleton::SkeletonSlot;
9use crate::tensor::skeleton::{Clean, Tainting};
10use crate::tensor::storage::TensorData;
11use crate::tensor::traits::{Composable, Dimension, Numeric, Operand};
12use std::ops::Index;
13use std::sync::Arc;
14
15/// Allocated tensor data exposed through the public API.
16///
17/// Internally, a `Tensor<T>` is an `Arc<TensorGraphEdge<T>>`: a
18/// reference-counted leaf node that carries the concrete `TensorData<T>`
19/// (data buffer + [`Layout`]) and the unique ID that the execution planner uses
20/// to track this value in the graph.
21///
22/// [`Layout`]: crate::tensor::mem_formats::layout::Layout
23///
24/// # Examples
25///
26/// ```
27/// use candela::Tensor;
28///
29/// // Fill every element with the same value.
30/// let t = Tensor::from_scalar(1.0_f64, &[3, 3]);
31/// assert_eq!(t.data().len(), 9);
32///
33/// // From an existing vec - total elements must equal the product of `shape`.
34/// let t = Tensor::from_vec(vec![1.0_f64, 2.0, 3.0], &[3]);
35/// assert_eq!(t.data(), &vec![1.0, 2.0, 3.0]);
36///
37/// // From any iterator.
38/// let t = Tensor::from_iter([1.0_f64, 2.0, 3.0, 4.0], &[4]);
39/// assert_eq!(t.data().len(), 4);
40///
41/// // Ops build a graph and run when you materialize.
42/// let t = Tensor::from_scalar(3.0_f64, &[4]);
43/// let result = (t * 2.0 + 1.0).materialize();
44/// assert_eq!(result.data(), &vec![7.0; 4]);
45/// ```
46pub struct Tensor<T, B: Backend = DefaultBackend> {
47 pub(crate) graph: Arc<TensorGraphEdge<T, B>>,
48}
49
50impl<T: ComputeFor<DefaultBackend>> Tensor<T> {
51 /// Create a tensor with every element set to `scalar`.
52 #[inline]
53 pub fn from_scalar(scalar: T, shape: &[usize]) -> Self {
54 Self {
55 graph: Arc::new(TensorGraphEdge::from_tensor_data(TensorData::from_scalar(
56 scalar, shape,
57 ))),
58 }
59 }
60
61 /// Create a tensor from `vector` interpreted with `shape`.
62 ///
63 /// # Panics
64 ///
65 /// Panics if `vector` length does not equal the product of `shape`.
66 #[inline]
67 pub fn from_vec(vector: Vec<T>, shape: &[usize]) -> Self {
68 Self {
69 graph: Arc::new(TensorGraphEdge::from_tensor_data(TensorData::from_vec(
70 vector, shape, 0,
71 ))),
72 }
73 }
74
75 /// Create a tensor by copying `data` into a buffer with the given `shape`.
76 ///
77 /// # Panics
78 ///
79 /// Panics if `data` length does not equal the product of `shape`.
80 #[inline]
81 pub fn from_slice(data: &[T], shape: &[usize]) -> Self {
82 Self::from_vec(data.to_vec(), shape)
83 }
84
85 /// Create a tensor by collecting `iter` into a buffer with the given `shape`.
86 ///
87 /// # Panics
88 ///
89 /// Panics if `iter` length does not equal the product of `shape`.
90 #[inline]
91 pub fn from_iter<I>(iter: I, shape: &[usize]) -> Self
92 where
93 I: IntoIterator<Item = T>,
94 {
95 let vector: Vec<T> = std::vec::Vec::from_iter(iter);
96 Self::from_vec(vector, shape)
97 }
98
99 /// Create an `n`×`m` matrix with ones on the main diagonal and zeros elsewhere.
100 ///
101 /// With `n == m` this is the identity matrix.
102 ///
103 /// # Examples
104 ///
105 /// ```
106 /// use candela::Tensor;
107 /// let i: Tensor<f64> = Tensor::eye(2, 2);
108 /// assert_eq!(i.data(), &[1.0, 0.0, 0.0, 1.0]);
109 /// ```
110 #[inline]
111 pub fn eye(n: usize, m: usize) -> Self {
112 let mut data: Vec<T> = vec![T::ZERO; n * m];
113
114 let mut i: usize = 0;
115 while i < data.len() {
116 data[i] = T::ONE;
117 i += m + 1;
118 }
119
120 Self::from_vec(data, &[n, m])
121 }
122}
123
124impl<T: Clone, B: Backend> Tensor<T, B> {
125 #[inline]
126 pub(crate) fn from_data(data: TensorData<T>) -> Self {
127 Self {
128 graph: Arc::new(TensorGraphEdge::from_tensor_data(data)),
129 }
130 }
131
132 /// Returns a reference to the underlying data buffer.
133 ///
134 /// Slicing, transposition, etc will change the layout of the tensor
135 /// so this is not guaranteed to be what you expect the tensor to
136 /// logically contain.
137 ///
138 /// Use [`.iter()`][Self::iter], to iterate over the whole tensor following
139 /// logical order or [`.index()`][Self::index] to access a single element by index.
140 ///
141 /// # Examples
142 ///
143 /// ```
144 /// use candela::Tensor;
145 /// let t = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]);
146 /// assert_eq!(t.data(), &[1.0, 2.0, 3.0, 4.0]);
147 /// ```
148 #[inline]
149 pub fn data(&self) -> &[T] {
150 self.graph.get().data()
151 }
152
153 /// Iterate over the tensor's elements in logical (row-major) order.
154 ///
155 /// Unlike [`.data()`][Self::data], this follows the tensor's layout, so a
156 /// sliced or transposed tensor yields its elements in the order its shape
157 /// implies.
158 ///
159 /// # Examples
160 ///
161 /// ```
162 /// use candela::Tensor;
163 /// let t = Tensor::from_slice(&[1.0, 2.0, 3.0], &[3]);
164 /// let collected: Vec<f64> = t.iter().copied().collect();
165 /// assert_eq!(collected, vec![1.0, 2.0, 3.0]);
166 /// ```
167 #[inline]
168 pub fn iter(&self) -> Iter<'_, T> {
169 self.graph.get().iter()
170 }
171
172 /// Iterate over the backing buffer using `layout` instead of this tensor's own
173 /// layout. Useful for traversals more exotic than the safe interface exposes.
174 ///
175 /// # Safety
176 ///
177 /// `layout` must be a valid transformation of this tensor's current layout -
178 /// every index it addresses must fall within the backing buffer. A layout
179 /// derived from this tensor's layout (a view, slice, transpose, or broadcast
180 /// of it) upholds this; an unrelated layout may read out of bounds and is
181 /// undefined behaviour.
182 #[inline]
183 pub unsafe fn iter_as_layout<'a>(&'a self, layout: &'a Layout) -> Iter<'a, T> {
184 unsafe { self.graph.get().iter_as_layout(layout) }
185 }
186
187 /// Walk the tensor depth-first, yielding a [`StepInfo`] for each element and
188 /// for each dimension boundary crossed along the way.
189 ///
190 /// Unlike [`.iter()`][Self::iter], which yields a flat stream of elements,
191 /// these events carry enough structure to reconstruct the tensor's nesting.
192 /// The walk follows the logical layout, so a sliced or transposed tensor is
193 /// visited in the order its shape implies. It rebuilds that order one index
194 /// at a time and is not intended for hot paths. The
195 /// [`Display`](std::fmt::Display) implementation is built on it.
196 ///
197 /// # Examples
198 ///
199 /// Regroup a flat buffer back into its rows - something [`.iter()`][Self::iter]
200 /// alone can't do, because it never signals where one row ends and the next
201 /// begins:
202 ///
203 /// ```
204 /// use candela::{StepInfo, Tensor};
205 ///
206 /// let t = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]);
207 ///
208 /// let mut rows: Vec<Vec<f64>> = Vec::new();
209 /// for step in t.informed_iter() {
210 /// match step {
211 /// // The innermost dimension (axis 1) opening means a new row starts.
212 /// StepInfo::EnterDimension(1) => rows.push(Vec::new()),
213 /// StepInfo::Value(v) => rows.last_mut().unwrap().push(v),
214 /// _ => {}
215 /// }
216 /// }
217 /// assert_eq!(rows, vec![vec![1.0, 2.0], vec![3.0, 4.0]]);
218 /// ```
219 #[inline]
220 pub fn informed_iter(&self) -> InformedIter<'_, T> {
221 self.graph.get().informed_iter()
222 }
223
224 /// Makes a deep copy of this tensor.
225 #[inline]
226 pub fn deep_clone(&self) -> Self {
227 let data = self.graph.get();
228
229 Self {
230 graph: Arc::new(TensorGraphEdge::from_tensor_data(data.deep_clone())),
231 }
232 }
233
234 /// Make a shallow copy of this tensor with a new graph identity.
235 ///
236 /// The underlying buffer is shared with the original, but the new tensor carries a fresh
237 /// graph ID. The planner treats it as an unrelated input - no connection is maintained to
238 /// any live promises that reference the original. Use [`Tensor::clone`] to preserve that
239 /// connection, or [`Tensor::deep_clone`] for a fully independent buffer.
240 ///
241 /// [`Tensor::clone`]: Tensor::clone
242 /// [`Tensor::deep_clone`]: Tensor::deep_clone
243 #[inline]
244 pub fn clone_detached(&self) -> Self {
245 let data = self.graph.get();
246
247 Self {
248 graph: Arc::new(TensorGraphEdge::from_tensor_data(data.clone())),
249 }
250 }
251}
252
253impl<T: Numeric, B: Backend> Tensor<T, B> {
254 /// Wrap this tensor as a [`TensorPromise`] without applying any transformation.
255 ///
256 /// The primary use case is initializing a mutable accumulator that will have ops
257 /// applied to it in a loop - as it needs a [`TensorPromise<T>`] on both sides
258 /// of the assignment:
259 ///
260 /// ```
261 /// use candela::arange;
262 /// let t = arange!(4); // [0.0, 1.0, 2.0, 3.0]
263 /// let mut p = t.to_promise();
264 /// for i in 0..5_u32 {
265 /// p += i as f64;
266 /// }
267 /// // each element gains 0+1+2+3+4 = 10
268 /// assert_eq!(p.materialize().data(), &[10.0, 11.0, 12.0, 13.0]);
269 /// ```
270 ///
271 /// [`TensorPromise<T>`]: crate::tensor::promise::TensorPromise
272 #[inline]
273 pub fn to_promise(&self) -> TensorPromise<T, B> {
274 unsafe {
275 TensorPromise::new(
276 super::ops::def_op::OpKind::NoOp,
277 [NodeKind::Edge(self.graph.clone())].into(),
278 )
279 .unwrap_unchecked()
280 }
281 }
282
283 /// Return a reference to the element at `index`, following the tensor's layout.
284 ///
285 /// # Errors
286 ///
287 /// Returns [`OpError::NotEnoughAxes`] if `index` doesn't have one entry per
288 /// axis, or [`OpError::IndexOutOfBounds`] if an index is past the end of its axis.
289 ///
290 /// # Examples
291 ///
292 /// ```
293 /// use candela::Tensor;
294 /// let t = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3]);
295 /// assert_eq!(*t.get(&[1, 2])?, 6.0);
296 /// assert!(t.get(&[2, 0]).is_err()); // row 2 is past the end
297 /// # Ok::<(), candela::OpError>(())
298 /// ```
299 // TODO: Add support for negative indexing
300 pub fn get(&self, index: &[usize]) -> Result<&T, OpError> {
301 self.graph.get().get(index)
302 }
303
304 /// Return a reference to the tensor's first element.
305 ///
306 /// Most useful for reading a one-element result, such as a full reduction like
307 /// [`.sum()`][Self::sum].
308 ///
309 /// # Examples
310 ///
311 /// ```
312 /// use candela::Tensor;
313 /// let t = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], &[2, 2]);
314 /// let total = t.sum().materialize();
315 /// assert_eq!(*total.item(), 10.0);
316 /// ```
317 pub fn item(&self) -> &T {
318 self.graph.get().item()
319 }
320
321 /// Creates a [`SkeletonSlot`] shaped like this tensor.
322 ///
323 /// The slot is an input placeholder for a [`Skeleton`]. It has the tensor's
324 /// [`Layout`] but holds no data.
325 ///
326 /// [`Skeleton`]: crate::skeleton::Skeleton
327 ///
328 /// # Examples
329 ///
330 /// ```
331 /// use candela::Tensor;
332 ///
333 /// let a = Tensor::from_scalar(0.3, &[4]);
334 /// let slot = a.to_slot(); // placeholder shaped like a
335 /// let skeleton = (&slot * 2.0 + 1.0).into_skeleton(&[slot])?;
336 /// assert!(skeleton.run(&[&a]).is_ok());
337 /// # Ok::<(), candela::OpError>(())
338 /// ```
339 pub fn to_slot(&self) -> SkeletonSlot<T, B> {
340 SkeletonSlot::new(self.layout().clone())
341 }
342}
343
344impl<T, B: Backend> Dimension for Tensor<T, B> {
345 #[inline]
346 fn layout(&self) -> &super::mem_formats::layout::Layout {
347 self.graph.layout()
348 }
349}
350
351impl<T, B: Backend> Operand<T, B> for Tensor<T, B> {
352 fn to_node(&self) -> NodeKind<T, B> {
353 NodeKind::Edge(self.graph.clone())
354 }
355}
356
357impl<T, B: Backend> Tainting for Tensor<T, B> {
358 type Mark = Clean;
359}
360
361impl<T, B: Backend> Composable<T, B> for Tensor<T, B> {}
362
363impl<T, B: Backend> Clone for Tensor<T, B> {
364 /// Shallow copy sharing the same underlying buffer and graph identity.
365 ///
366 /// Equivalent to bumping an `Arc` reference count. The copy is connected to all promises
367 /// that reference the original - the planner sees them as the same input node. For a copy
368 /// the graph treats as unrelated, use [`clone_detached`]. For an independent buffer, use
369 /// [`deep_clone`].
370 ///
371 /// [`clone_detached`]: Tensor::clone_detached
372 /// [`deep_clone`]: Tensor::deep_clone
373 #[inline]
374 fn clone(&self) -> Self {
375 Self {
376 graph: self.graph.clone(),
377 }
378 }
379}
380
381#[allow(private_bounds)]
382impl<T: std::fmt::Display + Copy, B: Backend> std::fmt::Debug for Tensor<T, B> {
383 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
384 writeln!(f, "Tensor {:?}", self.layout())?;
385 std::fmt::Display::fmt(self, f)
386 }
387}
388
389impl<T, B> Index<&[usize]> for Tensor<T, B>
390where
391 T: Copy,
392 B: Backend,
393{
394 type Output = T;
395
396 fn index(&self, index: &[usize]) -> &Self::Output {
397 &self.graph.get()[index]
398 }
399}
400
401#[allow(private_bounds)]
402impl<T: std::fmt::Display + Copy, B: Backend> std::fmt::Display for Tensor<T, B> {
403 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
404 let mut indent = 0;
405 let mut in_seq = false;
406
407 debug_assert!(!self.shape().is_empty(), "Tensor rank must be >= 1");
408 let last = self.shape().len() - 1;
409
410 for step in self.informed_iter() {
411 match step {
412 StepInfo::EnterDimension(dim) => {
413 write!(f, "{:indent$}[", "", indent = indent)?;
414 indent += 2;
415
416 if dim != last {
417 writeln!(f)?;
418 }
419 }
420 StepInfo::ExitDimension(dim) => {
421 indent -= 2;
422 in_seq = false;
423
424 if dim != last {
425 write!(f, "{:indent$}", "", indent = indent)?;
426 }
427
428 writeln!(f, "]")?;
429 }
430 StepInfo::Value(v) => {
431 if in_seq {
432 write!(f, ", ")?;
433 }
434
435 write!(f, "{:>4}", v)?;
436
437 in_seq = true;
438 }
439 _ => {}
440 }
441 }
442
443 Ok(())
444 }
445}