tenferro_ops/
std_tensor_op.rs

1use std::hash::{Hash, Hasher};
2use std::sync::Arc;
3
4#[cfg(feature = "autodiff")]
5use computegraph::types::{LocalValueId, OperationRole, ValueKey};
6use computegraph::GraphOperation;
7use num_complex::{Complex32, Complex64};
8#[cfg(feature = "autodiff")]
9use tidu::{ADRuleResult, Primitive, PrimitiveBuilder};
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::ReduceProd { axes: a }, Self::ReduceProd { axes: b })
185            | (Self::ReduceMax { axes: a }, Self::ReduceMax { axes: b })
186            | (Self::ReduceMin { axes: a }, Self::ReduceMin { axes: b })
187            | (Self::Reverse { axes: a }, Self::Reverse { axes: b }) => a == b,
188            (Self::Compare(a), Self::Compare(b)) => a == b,
189            (
190                Self::ExtractDiag {
191                    axis_a: aa,
192                    axis_b: ba,
193                },
194                Self::ExtractDiag {
195                    axis_a: ab,
196                    axis_b: bb,
197                },
198            )
199            | (
200                Self::EmbedDiag {
201                    axis_a: aa,
202                    axis_b: ba,
203                },
204                Self::EmbedDiag {
205                    axis_a: ab,
206                    axis_b: bb,
207                },
208            ) => aa == ab && ba == bb,
209            (Self::Tril { k: a }, Self::Tril { k: b })
210            | (Self::Triu { k: a }, Self::Triu { k: b }) => a == b,
211            (Self::Gather(a), Self::Gather(b)) => a == b,
212            (
213                Self::GatherDynamicSliceSizes {
214                    offset_dims: oa,
215                    collapsed_slice_dims: ca,
216                    start_index_map: sa,
217                    index_vector_dim: ia,
218                    slice_sizes: za,
219                },
220                Self::GatherDynamicSliceSizes {
221                    offset_dims: ob,
222                    collapsed_slice_dims: cb,
223                    start_index_map: sb,
224                    index_vector_dim: ib,
225                    slice_sizes: zb,
226                },
227            ) => oa == ob && ca == cb && sa == sb && ia == ib && za == zb,
228            (Self::Scatter(a), Self::Scatter(b)) => a == b,
229            (Self::Slice(a), Self::Slice(b)) => a == b,
230            (Self::DynamicSlice { slice_sizes: a }, Self::DynamicSlice { slice_sizes: b }) => {
231                a == b
232            }
233            (Self::Pad(a), Self::Pad(b)) => a == b,
234            (
235                Self::Concatenate {
236                    axis: a,
237                    input_count: na,
238                },
239                Self::Concatenate {
240                    axis: b,
241                    input_count: nb,
242                },
243            ) => a == b && na == nb,
244            (Self::ShapeOf { axis: a }, Self::ShapeOf { axis: b })
245            | (Self::DynamicTruncate { axis: a }, Self::DynamicTruncate { axis: b })
246            | (Self::PadToMatch { axis: a }, Self::PadToMatch { axis: b }) => a == b,
247            (Self::Extension(a), Self::Extension(b)) => ext_op_eq(a.as_ref(), b.as_ref()),
248            _ => false,
249        }
250    }
251}
252
253impl Eq for StdTensorOp {}
254
255impl Hash for StdTensorOp {
256    fn hash<H: Hasher>(&self, state: &mut H) {
257        std::mem::discriminant(self).hash(state);
258        match self {
259            Self::Add
260            | Self::Sub
261            | Self::Mul
262            | Self::Neg
263            | Self::Conj
264            | Self::Div
265            | Self::Rem
266            | Self::Abs
267            | Self::Sign
268            | Self::Maximum
269            | Self::Minimum
270            | Self::Select
271            | Self::Clamp
272            | Self::Exp
273            | Self::Log
274            | Self::Sin
275            | Self::Cos
276            | Self::Tanh
277            | Self::Sqrt
278            | Self::Rsqrt
279            | Self::Pow
280            | Self::Expm1
281            | Self::Log1p => {}
282            Self::DotGeneral { config } => {
283                config.hash(state);
284            }
285            Self::Transpose { perm } => perm.hash(state),
286            Self::Reshape { to_shape } => {
287                to_shape.hash(state);
288            }
289            Self::BroadcastInDim { shape, dims } => {
290                shape.hash(state);
291                dims.hash(state);
292            }
293            Self::Convert { from, to } => {
294                from.hash(state);
295                to.hash(state);
296            }
297            Self::Constant { dtype, bytes } => {
298                dtype.hash(state);
299                bytes.hash(state);
300            }
301            Self::ReduceSum { axes } => {
302                axes.hash(state);
303            }
304            Self::Compare(dir) => dir.hash(state),
305            Self::ExtractDiag { axis_a, axis_b } | Self::EmbedDiag { axis_a, axis_b } => {
306                axis_a.hash(state);
307                axis_b.hash(state);
308            }
309            Self::Tril { k } | Self::Triu { k } => k.hash(state),
310            Self::Gather(config) => config.hash(state),
311            Self::GatherDynamicSliceSizes {
312                offset_dims,
313                collapsed_slice_dims,
314                start_index_map,
315                index_vector_dim,
316                slice_sizes,
317            } => {
318                offset_dims.hash(state);
319                collapsed_slice_dims.hash(state);
320                start_index_map.hash(state);
321                index_vector_dim.hash(state);
322                slice_sizes.hash(state);
323            }
324            Self::Scatter(config) => config.hash(state),
325            Self::Slice(config) => config.hash(state),
326            Self::DynamicSlice { slice_sizes } => slice_sizes.hash(state),
327            Self::DynamicUpdateSlice => {}
328            Self::Pad(config) => config.hash(state),
329            Self::Concatenate { axis, input_count } => {
330                axis.hash(state);
331                input_count.hash(state);
332            }
333            Self::Reverse { axes } => axes.hash(state),
334            Self::ShapeOf { axis } | Self::DynamicTruncate { axis } | Self::PadToMatch { axis } => {
335                axis.hash(state)
336            }
337            Self::ReduceProd { axes } | Self::ReduceMax { axes } | Self::ReduceMin { axes } => {
338                axes.hash(state);
339            }
340            Self::Extension(op) => hash_extension(op.as_ref(), state),
341        }
342    }
343}
344
345fn n_inputs_from_dim_exprs(min_inputs: usize, exprs: &[&[DimExpr]]) -> usize {
346    let max_idx = exprs
347        .iter()
348        .flat_map(|exprs| exprs.iter())
349        .filter_map(DimExpr::max_input_idx)
350        .max()
351        .map_or(0, |max_idx| max_idx + 1);
352    max_idx.max(min_inputs)
353}
354
355impl GraphOperation for StdTensorOp {
356    type Operand = tenferro_tensor::Tensor;
357    type Context = ();
358    type InputKey = TensorInputKey;
359
360    fn input_count(&self) -> usize {
361        match self {
362            Self::Add | Self::Sub | Self::Mul | Self::DotGeneral { .. } | Self::Gather(_) => 2,
363            Self::GatherDynamicSliceSizes { slice_sizes, .. } => {
364                n_inputs_from_dim_exprs(2, &[slice_sizes])
365            }
366            Self::Neg
367            | Self::Conj
368            | Self::Transpose { .. }
369            | Self::Convert { .. }
370            | Self::ExtractDiag { .. }
371            | Self::EmbedDiag { .. }
372            | Self::Tril { .. }
373            | Self::Triu { .. }
374            | Self::Slice(_)
375            | Self::Pad(_)
376            | Self::Reverse { .. }
377            | Self::ShapeOf { .. } => 1,
378            Self::DynamicTruncate { .. } | Self::PadToMatch { .. } => 2,
379            Self::Reshape { to_shape } => n_inputs_from_dim_exprs(1, &[to_shape]),
380            Self::BroadcastInDim { shape, .. } => n_inputs_from_dim_exprs(1, &[shape]),
381            Self::ReduceSum { .. }
382            | Self::ReduceProd { .. }
383            | Self::ReduceMax { .. }
384            | Self::ReduceMin { .. } => 1,
385            Self::Div
386            | Self::Rem
387            | Self::Maximum
388            | Self::Minimum
389            | Self::Pow
390            | Self::DynamicSlice { .. } => 2,
391            Self::Constant { .. } => 0,
392            Self::Scatter(_) | Self::DynamicUpdateSlice => 3,
393            Self::Concatenate { input_count, .. } => *input_count,
394            Self::Abs
395            | Self::Sign
396            | Self::Exp
397            | Self::Log
398            | Self::Sin
399            | Self::Cos
400            | Self::Tanh
401            | Self::Sqrt
402            | Self::Rsqrt
403            | Self::Expm1
404            | Self::Log1p => 1,
405            Self::Select | Self::Clamp => 3,
406            Self::Compare(_) => 2,
407            Self::Extension(op) => ExtensionOp::input_count(op.as_ref()),
408        }
409    }
410
411    fn output_count(&self) -> usize {
412        match self {
413            Self::Add
414            | Self::Sub
415            | Self::Mul
416            | Self::Neg
417            | Self::Conj
418            | Self::DotGeneral { .. }
419            | Self::Transpose { .. }
420            | Self::Reshape { .. }
421            | Self::BroadcastInDim { .. }
422            | Self::Convert { .. }
423            | Self::ReduceSum { .. }
424            | Self::Div
425            | Self::Rem
426            | Self::Abs
427            | Self::Sign
428            | Self::Maximum
429            | Self::Minimum
430            | Self::Compare(_)
431            | Self::Select
432            | Self::Clamp
433            | Self::Constant { .. }
434            | Self::Exp
435            | Self::Log
436            | Self::Sin
437            | Self::Cos
438            | Self::Tanh
439            | Self::Sqrt
440            | Self::Rsqrt
441            | Self::Pow
442            | Self::Expm1
443            | Self::Log1p
444            | Self::ExtractDiag { .. }
445            | Self::EmbedDiag { .. }
446            | Self::Tril { .. }
447            | Self::Triu { .. }
448            | Self::Gather(_)
449            | Self::GatherDynamicSliceSizes { .. }
450            | Self::Scatter(_)
451            | Self::Slice(_)
452            | Self::DynamicSlice { .. }
453            | Self::DynamicUpdateSlice
454            | Self::Pad(_)
455            | Self::Reverse { .. }
456            | Self::ShapeOf { .. }
457            | Self::DynamicTruncate { .. }
458            | Self::PadToMatch { .. }
459            | Self::ReduceProd { .. }
460            | Self::ReduceMax { .. }
461            | Self::ReduceMin { .. } => 1,
462            Self::Concatenate { .. } => 1,
463            Self::Extension(op) => ExtensionOp::output_count(op.as_ref()),
464        }
465    }
466}
467
468#[cfg(feature = "autodiff")]
469impl Primitive for StdTensorOp {
470    type ADContext = crate::ad::context::ShapeGuardContext;
471
472    fn add() -> Self {
473        StdTensorOp::Add
474    }
475
476    fn jvp_rule(
477        &self,
478        builder: &mut impl PrimitiveBuilder<Self>,
479        primal_in: &[ValueKey<Self>],
480        primal_out: &[ValueKey<Self>],
481        tangent_in: &[Option<LocalValueId>],
482        ctx: &mut Self::ADContext,
483    ) -> ADRuleResult<Vec<Option<LocalValueId>>> {
484        crate::ad::linearize(self, builder, primal_in, primal_out, tangent_in, ctx)
485    }
486
487    fn transpose_rule(
488        &self,
489        builder: &mut impl PrimitiveBuilder<Self>,
490        cotangent_out: &[Option<LocalValueId>],
491        inputs: &[tidu::PrimitiveTransposeInput<Self>],
492        mode: &OperationRole,
493        ctx: &mut Self::ADContext,
494    ) -> ADRuleResult<Vec<Option<LocalValueId>>> {
495        crate::ad::transpose_rule(self, builder, cotangent_out, inputs, mode, ctx)
496    }
497}
498
499#[cfg(all(test, feature = "autodiff"))]
500impl StdTensorOp {
501    pub(crate) fn jvp_rule(
502        &self,
503        builder: &mut computegraph::graph::GraphBuilder<Self>,
504        primal_in: &[ValueKey<Self>],
505        primal_out: &[ValueKey<Self>],
506        tangent_in: &[Option<LocalValueId>],
507        ctx: &mut crate::ad::context::ShapeGuardContext,
508    ) -> ADRuleResult<Vec<Option<LocalValueId>>> {
509        crate::ad::linearize(self, builder, primal_in, primal_out, tangent_in, ctx)
510    }
511
512    pub(crate) fn transpose_rule(
513        &self,
514        builder: &mut computegraph::graph::GraphBuilder<Self>,
515        cotangent_out: &[Option<LocalValueId>],
516        inputs: &[computegraph::ValueRef<Self>],
517        mode: &OperationRole,
518        ctx: &mut crate::ad::context::ShapeGuardContext,
519    ) -> ADRuleResult<Vec<Option<LocalValueId>>> {
520        let inputs = inputs
521            .iter()
522            .map(|input| match input {
523                computegraph::ValueRef::Local(local_id) => {
524                    let key = builder.global_key(*local_id).clone();
525                    tidu::PrimitiveTransposeInput::Residual(key)
526                }
527                computegraph::ValueRef::External(key) => {
528                    tidu::PrimitiveTransposeInput::Residual(key.clone())
529                }
530            })
531            .collect::<Vec<_>>();
532        crate::ad::transpose_rule(self, builder, cotangent_out, inputs.as_slice(), mode, ctx)
533    }
534}
tenferro_ops/std_tensor_op.rs

tenferro_ops/
std_tensor_op.rs
