tenferro_ops/
std_tensor_op.rs

1use std::hash::{Hash, Hasher};
2use std::sync::Arc;
3
4#[cfg(all(test, feature = "autodiff"))]
5use crate::ad::{ADRuleResult, PrimitiveTransposeInput};
6#[cfg(all(test, feature = "autodiff"))]
7use computegraph::types::{LocalValueId, OperationRole, ValueKey};
8use computegraph::GraphOperation;
9use num_complex::{Complex32, Complex64};
10
11use crate::dim_expr::DimExpr;
12use crate::ext_op::{ext_op_eq, hash_extension, ExtensionOp};
13use crate::input_key::TensorInputKey;
14use tenferro_tensor::{
15    CompareDir, DType, DotGeneralConfig, GatherConfig, PadConfig, ScatterConfig, SliceConfig,
16    TensorScalar,
17};
18
19/// Scalar values that can be encoded as tensor constant operations.
20///
21/// # Examples
22///
23/// ```rust
24/// use tenferro_ops::std_tensor_op::ConstantScalar;
25///
26/// assert_eq!(1.0_f64.constant_bytes(), 1.0_f64.to_le_bytes().to_vec());
27/// ```
28pub trait ConstantScalar: TensorScalar + private::Sealed {
29    /// Encode the scalar value as little-endian constant bytes.
30    ///
31    /// # Examples
32    ///
33    /// ```rust
34    /// use tenferro_ops::std_tensor_op::ConstantScalar;
35    ///
36    /// assert_eq!(true.constant_bytes(), vec![1]);
37    /// ```
38    fn constant_bytes(self) -> Vec<u8>;
39}
40
41mod private {
42    pub trait Sealed {}
43
44    impl Sealed for f64 {}
45    impl Sealed for f32 {}
46    impl Sealed for i64 {}
47    impl Sealed for i32 {}
48    impl Sealed for bool {}
49    impl Sealed for num_complex::Complex64 {}
50    impl Sealed for num_complex::Complex32 {}
51}
52
53impl ConstantScalar for f64 {
54    fn constant_bytes(self) -> Vec<u8> {
55        self.to_le_bytes().to_vec()
56    }
57}
58
59impl ConstantScalar for f32 {
60    fn constant_bytes(self) -> Vec<u8> {
61        self.to_le_bytes().to_vec()
62    }
63}
64
65impl ConstantScalar for i64 {
66    fn constant_bytes(self) -> Vec<u8> {
67        self.to_le_bytes().to_vec()
68    }
69}
70
71impl ConstantScalar for i32 {
72    fn constant_bytes(self) -> Vec<u8> {
73        self.to_le_bytes().to_vec()
74    }
75}
76
77impl ConstantScalar for bool {
78    fn constant_bytes(self) -> Vec<u8> {
79        vec![u8::from(self)]
80    }
81}
82
83impl ConstantScalar for Complex64 {
84    fn constant_bytes(self) -> Vec<u8> {
85        let mut bytes = Vec::with_capacity(16);
86        bytes.extend_from_slice(&self.re.to_le_bytes());
87        bytes.extend_from_slice(&self.im.to_le_bytes());
88        bytes
89    }
90}
91
92impl ConstantScalar for Complex32 {
93    fn constant_bytes(self) -> Vec<u8> {
94        let mut bytes = Vec::with_capacity(8);
95        bytes.extend_from_slice(&self.re.to_le_bytes());
96        bytes.extend_from_slice(&self.im.to_le_bytes());
97        bytes
98    }
99}
100
101tenferro_core_ops::define_std_tensor_op!();
102
103impl StdTensorOp {
104    /// Create a scalar constant op from any supported tensor scalar.
105    ///
106    /// # Examples
107    ///
108    /// ```rust
109    /// use num_complex::Complex64;
110    /// use tenferro_ops::std_tensor_op::StdTensorOp;
111    /// use tenferro_tensor::DType;
112    ///
113    /// let real = StdTensorOp::constant(1.5_f64);
114    /// let complex = StdTensorOp::constant(Complex64::new(1.0, -2.0));
115    ///
116    /// assert!(matches!(real, StdTensorOp::Constant { dtype: DType::F64, .. }));
117    /// assert!(matches!(complex, StdTensorOp::Constant { dtype: DType::C64, .. }));
118    /// ```
119    pub fn constant<T: ConstantScalar>(value: T) -> Self {
120        Self::Constant {
121            dtype: T::dtype(),
122            bytes: value.constant_bytes(),
123        }
124    }
125}
126
127impl PartialEq for StdTensorOp {
128    fn eq(&self, other: &Self) -> bool {
129        if std::mem::discriminant(self) != std::mem::discriminant(other) {
130            return false;
131        }
132        match (self, other) {
133            (Self::Add, Self::Add)
134            | (Self::Sub, Self::Sub)
135            | (Self::Mul, Self::Mul)
136            | (Self::Neg, Self::Neg)
137            | (Self::Conj, Self::Conj)
138            | (Self::Div, Self::Div)
139            | (Self::Rem, Self::Rem)
140            | (Self::Abs, Self::Abs)
141            | (Self::Sign, Self::Sign)
142            | (Self::Maximum, Self::Maximum)
143            | (Self::Minimum, Self::Minimum)
144            | (Self::Select, Self::Select)
145            | (Self::Clamp, Self::Clamp)
146            | (Self::Exp, Self::Exp)
147            | (Self::Log, Self::Log)
148            | (Self::Sin, Self::Sin)
149            | (Self::Cos, Self::Cos)
150            | (Self::Tanh, Self::Tanh)
151            | (Self::Sqrt, Self::Sqrt)
152            | (Self::Rsqrt, Self::Rsqrt)
153            | (Self::Pow, Self::Pow)
154            | (Self::Expm1, Self::Expm1)
155            | (Self::Log1p, Self::Log1p)
156            | (Self::DynamicUpdateSlice, Self::DynamicUpdateSlice) => true,
157            (Self::DotGeneral { config: a }, Self::DotGeneral { config: b }) => a == b,
158            (Self::Transpose { perm: a }, Self::Transpose { perm: b }) => a == b,
159            (Self::Reshape { to_shape: a }, Self::Reshape { to_shape: b }) => a == b,
160            (
161                Self::BroadcastInDim {
162                    shape: sa,
163                    dims: da,
164                },
165                Self::BroadcastInDim {
166                    shape: sb,
167                    dims: db,
168                },
169            ) => sa == sb && da == db,
170            (Self::Convert { from: fa, to: ta }, Self::Convert { from: fb, to: tb }) => {
171                fa == fb && ta == tb
172            }
173            (
174                Self::Constant {
175                    dtype: da,
176                    bytes: ba,
177                },
178                Self::Constant {
179                    dtype: db,
180                    bytes: bb,
181                },
182            ) => da == db && ba == bb,
183            (Self::ReduceSum { axes: a }, Self::ReduceSum { axes: b })
184            | (Self::ReduceSumSquares { axes: a }, Self::ReduceSumSquares { axes: b })
185            | (Self::ReduceProd { axes: a }, Self::ReduceProd { axes: b })
186            | (Self::ReduceMax { axes: a }, Self::ReduceMax { axes: b })
187            | (Self::ReduceMin { axes: a }, Self::ReduceMin { axes: b })
188            | (Self::Reverse { axes: a }, Self::Reverse { axes: b }) => a == b,
189            (Self::Compare(a), Self::Compare(b)) => a == b,
190            (
191                Self::ExtractDiag {
192                    axis_a: aa,
193                    axis_b: ba,
194                },
195                Self::ExtractDiag {
196                    axis_a: ab,
197                    axis_b: bb,
198                },
199            )
200            | (
201                Self::EmbedDiag {
202                    axis_a: aa,
203                    axis_b: ba,
204                },
205                Self::EmbedDiag {
206                    axis_a: ab,
207                    axis_b: bb,
208                },
209            ) => aa == ab && ba == bb,
210            (Self::Tril { k: a }, Self::Tril { k: b })
211            | (Self::Triu { k: a }, Self::Triu { k: b }) => a == b,
212            (Self::Gather(a), Self::Gather(b)) => a == b,
213            (
214                Self::GatherDynamicSliceSizes {
215                    offset_dims: oa,
216                    collapsed_slice_dims: ca,
217                    start_index_map: sa,
218                    index_vector_dim: ia,
219                    slice_sizes: za,
220                },
221                Self::GatherDynamicSliceSizes {
222                    offset_dims: ob,
223                    collapsed_slice_dims: cb,
224                    start_index_map: sb,
225                    index_vector_dim: ib,
226                    slice_sizes: zb,
227                },
228            ) => oa == ob && ca == cb && sa == sb && ia == ib && za == zb,
229            (Self::Scatter(a), Self::Scatter(b)) => a == b,
230            (Self::Slice(a), Self::Slice(b)) => a == b,
231            (Self::DynamicSlice { slice_sizes: a }, Self::DynamicSlice { slice_sizes: b }) => {
232                a == b
233            }
234            (Self::Pad(a), Self::Pad(b)) => a == b,
235            (
236                Self::Concatenate {
237                    axis: a,
238                    input_count: na,
239                },
240                Self::Concatenate {
241                    axis: b,
242                    input_count: nb,
243                },
244            ) => a == b && na == nb,
245            (Self::ShapeOf { axis: a }, Self::ShapeOf { axis: b })
246            | (Self::DynamicTruncate { axis: a }, Self::DynamicTruncate { axis: b })
247            | (Self::PadToMatch { axis: a }, Self::PadToMatch { axis: b }) => a == b,
248            (Self::Extension(a), Self::Extension(b)) => ext_op_eq(a.as_ref(), b.as_ref()),
249            _ => false,
250        }
251    }
252}
253
254impl Eq for StdTensorOp {}
255
256impl Hash for StdTensorOp {
257    fn hash<H: Hasher>(&self, state: &mut H) {
258        std::mem::discriminant(self).hash(state);
259        match self {
260            Self::Add
261            | Self::Sub
262            | Self::Mul
263            | Self::Neg
264            | Self::Conj
265            | Self::Div
266            | Self::Rem
267            | Self::Abs
268            | Self::Sign
269            | Self::Maximum
270            | Self::Minimum
271            | Self::Select
272            | Self::Clamp
273            | Self::Exp
274            | Self::Log
275            | Self::Sin
276            | Self::Cos
277            | Self::Tanh
278            | Self::Sqrt
279            | Self::Rsqrt
280            | Self::Pow
281            | Self::Expm1
282            | Self::Log1p => {}
283            Self::DotGeneral { config } => {
284                config.hash(state);
285            }
286            Self::Transpose { perm } => perm.hash(state),
287            Self::Reshape { to_shape } => {
288                to_shape.hash(state);
289            }
290            Self::BroadcastInDim { shape, dims } => {
291                shape.hash(state);
292                dims.hash(state);
293            }
294            Self::Convert { from, to } => {
295                from.hash(state);
296                to.hash(state);
297            }
298            Self::Constant { dtype, bytes } => {
299                dtype.hash(state);
300                bytes.hash(state);
301            }
302            Self::ReduceSum { axes } | Self::ReduceSumSquares { axes } => {
303                axes.hash(state);
304            }
305            Self::Compare(dir) => dir.hash(state),
306            Self::ExtractDiag { axis_a, axis_b } | Self::EmbedDiag { axis_a, axis_b } => {
307                axis_a.hash(state);
308                axis_b.hash(state);
309            }
310            Self::Tril { k } | Self::Triu { k } => k.hash(state),
311            Self::Gather(config) => config.hash(state),
312            Self::GatherDynamicSliceSizes {
313                offset_dims,
314                collapsed_slice_dims,
315                start_index_map,
316                index_vector_dim,
317                slice_sizes,
318            } => {
319                offset_dims.hash(state);
320                collapsed_slice_dims.hash(state);
321                start_index_map.hash(state);
322                index_vector_dim.hash(state);
323                slice_sizes.hash(state);
324            }
325            Self::Scatter(config) => config.hash(state),
326            Self::Slice(config) => config.hash(state),
327            Self::DynamicSlice { slice_sizes } => slice_sizes.hash(state),
328            Self::DynamicUpdateSlice => {}
329            Self::Pad(config) => config.hash(state),
330            Self::Concatenate { axis, input_count } => {
331                axis.hash(state);
332                input_count.hash(state);
333            }
334            Self::Reverse { axes } => axes.hash(state),
335            Self::ShapeOf { axis } | Self::DynamicTruncate { axis } | Self::PadToMatch { axis } => {
336                axis.hash(state)
337            }
338            Self::ReduceProd { axes } | Self::ReduceMax { axes } | Self::ReduceMin { axes } => {
339                axes.hash(state);
340            }
341            Self::Extension(op) => hash_extension(op.as_ref(), state),
342        }
343    }
344}
345
346fn n_inputs_from_dim_exprs(min_inputs: usize, exprs: &[&[DimExpr]]) -> usize {
347    let max_idx = exprs
348        .iter()
349        .flat_map(|exprs| exprs.iter())
350        .filter_map(DimExpr::max_input_idx)
351        .max()
352        .map_or(0, |max_idx| max_idx + 1);
353    max_idx.max(min_inputs)
354}
355
356impl GraphOperation for StdTensorOp {
357    type Operand = tenferro_tensor::Tensor;
358    type Context = ();
359    type InputKey = TensorInputKey;
360
361    fn input_count(&self) -> usize {
362        match self {
363            Self::Add | Self::Sub | Self::Mul | Self::DotGeneral { .. } | Self::Gather(_) => 2,
364            Self::GatherDynamicSliceSizes { slice_sizes, .. } => {
365                n_inputs_from_dim_exprs(2, &[slice_sizes])
366            }
367            Self::Neg
368            | Self::Conj
369            | Self::Transpose { .. }
370            | Self::Convert { .. }
371            | Self::ExtractDiag { .. }
372            | Self::EmbedDiag { .. }
373            | Self::Tril { .. }
374            | Self::Triu { .. }
375            | Self::Slice(_)
376            | Self::Pad(_)
377            | Self::Reverse { .. }
378            | Self::ShapeOf { .. } => 1,
379            Self::DynamicTruncate { .. } | Self::PadToMatch { .. } => 2,
380            Self::Reshape { to_shape } => n_inputs_from_dim_exprs(1, &[to_shape]),
381            Self::BroadcastInDim { shape, .. } => n_inputs_from_dim_exprs(1, &[shape]),
382            Self::ReduceSum { .. }
383            | Self::ReduceSumSquares { .. }
384            | Self::ReduceProd { .. }
385            | Self::ReduceMax { .. }
386            | Self::ReduceMin { .. } => 1,
387            Self::Div
388            | Self::Rem
389            | Self::Maximum
390            | Self::Minimum
391            | Self::Pow
392            | Self::DynamicSlice { .. } => 2,
393            Self::Constant { .. } => 0,
394            Self::Scatter(_) | Self::DynamicUpdateSlice => 3,
395            Self::Concatenate { input_count, .. } => *input_count,
396            Self::Abs
397            | Self::Sign
398            | Self::Exp
399            | Self::Log
400            | Self::Sin
401            | Self::Cos
402            | Self::Tanh
403            | Self::Sqrt
404            | Self::Rsqrt
405            | Self::Expm1
406            | Self::Log1p => 1,
407            Self::Select | Self::Clamp => 3,
408            Self::Compare(_) => 2,
409            Self::Extension(op) => ExtensionOp::input_count(op.as_ref()),
410        }
411    }
412
413    fn output_count(&self) -> usize {
414        match self {
415            Self::Add
416            | Self::Sub
417            | Self::Mul
418            | Self::Neg
419            | Self::Conj
420            | Self::DotGeneral { .. }
421            | Self::Transpose { .. }
422            | Self::Reshape { .. }
423            | Self::BroadcastInDim { .. }
424            | Self::Convert { .. }
425            | Self::ReduceSum { .. }
426            | Self::ReduceSumSquares { .. }
427            | Self::Div
428            | Self::Rem
429            | Self::Abs
430            | Self::Sign
431            | Self::Maximum
432            | Self::Minimum
433            | Self::Compare(_)
434            | Self::Select
435            | Self::Clamp
436            | Self::Constant { .. }
437            | Self::Exp
438            | Self::Log
439            | Self::Sin
440            | Self::Cos
441            | Self::Tanh
442            | Self::Sqrt
443            | Self::Rsqrt
444            | Self::Pow
445            | Self::Expm1
446            | Self::Log1p
447            | Self::ExtractDiag { .. }
448            | Self::EmbedDiag { .. }
449            | Self::Tril { .. }
450            | Self::Triu { .. }
451            | Self::Gather(_)
452            | Self::GatherDynamicSliceSizes { .. }
453            | Self::Scatter(_)
454            | Self::Slice(_)
455            | Self::DynamicSlice { .. }
456            | Self::DynamicUpdateSlice
457            | Self::Pad(_)
458            | Self::Reverse { .. }
459            | Self::ShapeOf { .. }
460            | Self::DynamicTruncate { .. }
461            | Self::PadToMatch { .. }
462            | Self::ReduceProd { .. }
463            | Self::ReduceMax { .. }
464            | Self::ReduceMin { .. } => 1,
465            Self::Concatenate { .. } => 1,
466            Self::Extension(op) => ExtensionOp::output_count(op.as_ref()),
467        }
468    }
469}
470
471#[cfg(all(test, feature = "autodiff"))]
472impl StdTensorOp {
473    pub(crate) fn jvp_rule(
474        &self,
475        builder: &mut computegraph::graph::GraphBuilder<Self>,
476        primal_in: &[ValueKey<Self>],
477        primal_out: &[ValueKey<Self>],
478        tangent_in: &[Option<LocalValueId>],
479        ctx: &mut crate::ad::context::ShapeGuardContext,
480    ) -> ADRuleResult<Vec<Option<LocalValueId>>> {
481        crate::ad::linearize(self, builder, primal_in, primal_out, tangent_in, ctx)
482    }
483
484    pub(crate) fn transpose_rule(
485        &self,
486        builder: &mut computegraph::graph::GraphBuilder<Self>,
487        cotangent_out: &[Option<LocalValueId>],
488        inputs: &[computegraph::ValueRef<Self>],
489        mode: &OperationRole,
490        ctx: &mut crate::ad::context::ShapeGuardContext,
491    ) -> ADRuleResult<Vec<Option<LocalValueId>>> {
492        let inputs = inputs
493            .iter()
494            .map(|input| match input {
495                computegraph::ValueRef::Local(local_id) => {
496                    let key = builder.global_key(*local_id).clone();
497                    PrimitiveTransposeInput::Residual(key)
498                }
499                computegraph::ValueRef::External(key) => {
500                    PrimitiveTransposeInput::Residual(key.clone())
501                }
502            })
503            .collect::<Vec<_>>();
504        crate::ad::transpose_rule(self, builder, cotangent_out, inputs.as_slice(), mode, ctx)
505    }
506}
tenferro_ops/std_tensor_op.rs

tenferro_ops/
std_tensor_op.rs
