Skip to main content

candela/tensor/
traits.rs

1use crate::tensor::backend::Backend;
2use crate::tensor::graph::NodeKind;
3use crate::tensor::mem_formats::layout::Layout;
4use crate::tensor::storage::TensorData;
5
6/// Shape, stride, and layout queries shared by every tensor-like value.
7///
8/// Implemented by [`Tensor`](crate::Tensor), [`TensorPromise`](crate::TensorPromise),
9/// [`CachedTensorPromise`](crate::CachedTensorPromise), and skeleton slots, so
10/// `shape`, `len`, `is_contiguous`, and friends read the same on all of them -
11/// whether or not the value has been computed yet.
12///
13/// # Examples
14///
15/// ```
16/// use candela::{Dimension, Tensor};
17///
18/// fn element_count<D: Dimension>(x: &D) -> usize {
19///     x.len()
20/// }
21///
22/// let t = Tensor::from_scalar(1.0_f64, &[2, 3]);
23/// assert_eq!(t.shape(), &[2, 3]);
24/// assert_eq!(element_count(&t), 6);
25///
26/// // Works on an unevaluated promise too - the shape is known before compute.
27/// let p = &t + 1.0;
28/// assert_eq!(element_count(&p), 6);
29/// ```
30pub trait Dimension {
31    fn layout(&self) -> &Layout;
32
33    fn shape(&self) -> &'_ [usize] {
34        self.layout().shape()
35    }
36
37    fn stride(&self) -> &'_ [i32] {
38        self.layout().stride()
39    }
40
41    fn adj_stride(&self) -> &'_ [i32] {
42        self.layout().adj_stride()
43    }
44
45    fn len(&self) -> usize {
46        self.layout().len()
47    }
48
49    fn is_empty(&self) -> bool {
50        self.layout().is_empty()
51    }
52
53    fn offset(&self) -> usize {
54        self.layout().offset()
55    }
56
57    fn is_contiguous(&self) -> bool {
58        self.layout().is_contiguous()
59    }
60
61    fn is_contiguous_at_axis(&self, axis: usize) -> bool {
62        self.layout().is_contiguous_at_axis(axis)
63    }
64
65    fn is_transposed(&self) -> bool {
66        self.layout().is_transposed()
67    }
68
69    fn is_transposed_at_axis(&self, axis: usize) -> bool {
70        self.layout().is_transposed_at_axis(axis)
71    }
72}
73
74pub trait Promising {
75    type Output;
76
77    fn compute(&self) -> TensorData<Self::Output>;
78}
79
80/// Represents any graph type that can take part in an op: it produces a graph node and
81/// carries a layout. Implemented by every operand kind - `Tensor`,
82/// `TensorPromise`, `CachedTensorPromise`, `BakedPromise`, `SkeletonSlot`,
83/// `SkeletonPromise` - plus internal intermediates.
84pub(crate) trait Operand<T, B: Backend>: Dimension {
85    fn to_node(&self) -> NodeKind<T, B>;
86}
87
88/// A subset of all tensor types that can be materialized
89///
90/// A "materializable" subset of `Operand`: values that are legal as
91/// concrete inputs when binding a skeleton via `compose`. `SkeletonSlot` and
92/// `SkeletonPromise` are `Operand`s but not `Composable`. This
93/// exclusion guarantees that no unbound slot appears in the materialization
94/// path.
95///
96/// # Examples
97///
98/// ```
99/// use candela::skeleton::SkeletonSlot;
100/// use candela::Tensor;
101///
102/// let x = SkeletonSlot::from_shape(&[4]);
103/// let sk = (&x * 2.0).into_skeleton(&[x])?;
104///
105/// // `compose` accepts any `Composable` - a Tensor or a TensorPromise - but not a slot.
106/// let as_tensor = Tensor::from_scalar(3.0, &[4]);
107/// let as_promise = Tensor::from_scalar(3.0, &[4]) + 1.0;
108/// sk.compose(&[&as_tensor])?;
109/// sk.compose(&[&as_promise])?;
110/// # Ok::<(), candela::OpError>(())
111/// ```
112pub trait Composable<T, B: Backend>: Operand<T, B> {}
113
114pub trait StreamingIterator {
115    type Item<'a>
116    where
117        Self: 'a;
118
119    fn next_stream<'a>(&'a mut self) -> Option<Self::Item<'a>>;
120
121    #[allow(unused)]
122    fn zip<Other>(self, other: Other) -> StreamingZip<Self, Other>
123    where
124        Self: Sized,
125        Other: StreamingIterator,
126    {
127        StreamingZip {
128            left: self,
129            right: other,
130        }
131    }
132}
133
134#[allow(unused)]
135pub struct StreamingZip<A: StreamingIterator, B: StreamingIterator> {
136    left: A,
137    right: B,
138}
139
140impl<A, B> StreamingIterator for StreamingZip<A, B>
141where
142    A: StreamingIterator,
143    B: StreamingIterator,
144{
145    type Item<'a>
146        = (A::Item<'a>, B::Item<'a>)
147    where
148        Self: 'a;
149
150    fn next_stream<'a>(&'a mut self) -> Option<Self::Item<'a>> {
151        let l = self.left.next_stream()?;
152        let r = self.right.next_stream()?;
153        Some((l, r))
154    }
155}
156
157pub(crate) trait Numeric: crate::tensor::definitions::NumberLike {
158    const MUL_NEUTRAL: Self;
159    const SUM_NEUTRAL: Self;
160    const ONE: Self;
161    const ZERO: Self;
162    const MIN: Self;
163}
164
165impl Numeric for f64 {
166    const MUL_NEUTRAL: Self = 1.0;
167    const SUM_NEUTRAL: Self = 0.0;
168    const ONE: Self = 1.0;
169    const ZERO: Self = 0.0;
170    const MIN: Self = f64::NEG_INFINITY;
171}
172
173impl Numeric for f32 {
174    const MUL_NEUTRAL: Self = 1.0;
175    const SUM_NEUTRAL: Self = 0.0;
176    const ONE: Self = 1.0;
177    const ZERO: Self = 0.0;
178    const MIN: Self = f32::NEG_INFINITY;
179}
180
181pub(crate) trait FromIndex {
182    fn from_index(i: usize) -> Self;
183}
184
185impl FromIndex for f64 {
186    #[inline]
187    fn from_index(i: usize) -> Self {
188        i as f64
189    }
190}
191
192impl FromIndex for f32 {
193    #[inline]
194    fn from_index(i: usize) -> Self {
195        i as f32
196    }
197}