tenferro_core_ops/
catalog.rs

1/// High-level category for a core primitive operation.
2///
3/// # Examples
4///
5/// ```rust
6/// use tenferro_core_ops::{descriptor, OpCategory, PrimitiveOpKind};
7///
8/// assert_eq!(
9///     descriptor(PrimitiveOpKind::ShapeOf).category,
10///     OpCategory::Host
11/// );
12/// ```
13#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
14pub enum OpCategory {
15    Elementwise,
16    Analytic,
17    Structural,
18    Reduction,
19    Contraction,
20    Indexing,
21    Dynamic,
22    Host,
23}
24
25/// Dtype compatibility policy for a core primitive operation.
26///
27/// # Examples
28///
29/// ```rust
30/// use tenferro_core_ops::{descriptor, DTypePolicy, PrimitiveOpKind};
31///
32/// assert_eq!(
33///     descriptor(PrimitiveOpKind::Compare).dtype_policy,
34///     DTypePolicy::CompareToBool
35/// );
36/// ```
37#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
38pub enum DTypePolicy {
39    SameAny,
40    SameNumeric,
41    SameFloat,
42    /// Preserve real numeric dtype and map complex magnitude to the matching real dtype.
43    AbsToReal,
44    SameFloatOrComplex,
45    CompareToBool,
46    BoolSelect,
47    Convert,
48    Shape,
49    Constant,
50}
51
52/// Static metadata for one core primitive operation.
53///
54/// # Examples
55///
56/// ```rust
57/// use tenferro_core_ops::{descriptor, PrimitiveOpKind};
58///
59/// let add = descriptor(PrimitiveOpKind::Add);
60/// assert_eq!(add.min_inputs, 2);
61/// assert_eq!(add.max_inputs, 2);
62/// ```
63#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
64pub struct PrimitiveOpDescriptor {
65    /// Catalog key for this operation.
66    pub kind: PrimitiveOpKind,
67    /// Stable snake-case operation name for diagnostics and descriptors.
68    pub name: &'static str,
69    /// Broad execution category.
70    pub category: OpCategory,
71    /// Dtype compatibility policy.
72    pub dtype_policy: DTypePolicy,
73    /// Minimum number of inputs accepted by the op.
74    pub min_inputs: u8,
75    /// Maximum number of inputs accepted by the op.
76    pub max_inputs: u8,
77    /// Whether this op is executed by host/runtime logic rather than a tensor backend.
78    pub host_only: bool,
79}
80
81macro_rules! primitive_ops {
82    ($macro:ident) => {
83        $macro! {
84            Add, "add", Elementwise, SameNumeric, 2, 2, false;
85            Sub, "sub", Elementwise, SameNumeric, 2, 2, false;
86            Mul, "mul", Elementwise, SameNumeric, 2, 2, false;
87            Neg, "neg", Elementwise, SameNumeric, 1, 1, false;
88            Conj, "conj", Elementwise, SameFloatOrComplex, 1, 1, false;
89            Div, "div", Elementwise, SameNumeric, 2, 2, false;
90            Rem, "rem", Elementwise, SameNumeric, 2, 2, false;
91            Abs, "abs", Elementwise, AbsToReal, 1, 1, false;
92            Sign, "sign", Elementwise, SameNumeric, 1, 1, false;
93            Maximum, "maximum", Elementwise, SameNumeric, 2, 2, false;
94            Minimum, "minimum", Elementwise, SameNumeric, 2, 2, false;
95            Compare, "compare", Elementwise, CompareToBool, 2, 2, false;
96            Select, "select", Elementwise, BoolSelect, 3, 3, false;
97            Clamp, "clamp", Elementwise, SameFloat, 3, 3, false;
98            Exp, "exp", Analytic, SameFloatOrComplex, 1, 1, false;
99            Log, "log", Analytic, SameFloatOrComplex, 1, 1, false;
100            Sin, "sin", Analytic, SameFloatOrComplex, 1, 1, false;
101            Cos, "cos", Analytic, SameFloatOrComplex, 1, 1, false;
102            Tanh, "tanh", Analytic, SameFloatOrComplex, 1, 1, false;
103            Sqrt, "sqrt", Analytic, SameFloatOrComplex, 1, 1, false;
104            Rsqrt, "rsqrt", Analytic, SameFloatOrComplex, 1, 1, false;
105            Pow, "pow", Analytic, SameNumeric, 2, 2, false;
106            Expm1, "expm1", Analytic, SameFloatOrComplex, 1, 1, false;
107            Log1p, "log1p", Analytic, SameFloatOrComplex, 1, 1, false;
108            DotGeneral, "dot_general", Contraction, SameFloatOrComplex, 2, 2, false;
109            ReduceSum, "reduce_sum", Reduction, SameNumeric, 1, 1, false;
110            ReduceSumSquares, "reduce_sum_squares", Reduction, SameFloat, 1, 1, false;
111            ReduceProd, "reduce_prod", Reduction, SameNumeric, 1, 1, false;
112            ReduceMax, "reduce_max", Reduction, SameNumeric, 1, 1, false;
113            ReduceMin, "reduce_min", Reduction, SameNumeric, 1, 1, false;
114            Transpose, "transpose", Structural, SameAny, 1, 1, false;
115            Reshape, "reshape", Structural, SameAny, 1, 1, false;
116            BroadcastInDim, "broadcast_in_dim", Structural, SameAny, 1, 1, false;
117            Convert, "convert", Structural, Convert, 1, 1, false;
118            ExtractDiag, "extract_diag", Structural, SameAny, 1, 1, false;
119            EmbedDiag, "embed_diag", Structural, SameAny, 1, 1, false;
120            Tril, "tril", Structural, SameAny, 1, 1, false;
121            Triu, "triu", Structural, SameAny, 1, 1, false;
122            Gather, "gather", Indexing, SameAny, 2, 2, false;
123            GatherDynamicSliceSizes, "gather_dynamic_slice_sizes", Indexing, SameAny, 2, 2, false;
124            Scatter, "scatter", Indexing, SameAny, 3, 3, false;
125            Slice, "slice", Indexing, SameAny, 1, 1, false;
126            DynamicSlice, "dynamic_slice", Indexing, SameAny, 2, 2, false;
127            DynamicUpdateSlice, "dynamic_update_slice", Indexing, SameAny, 3, 3, false;
128            Pad, "pad", Indexing, SameAny, 1, 1, false;
129            Concatenate, "concatenate", Indexing, SameAny, 1, u8::MAX, false;
130            Reverse, "reverse", Indexing, SameAny, 1, 1, false;
131            ShapeOf, "shape_of", Host, Shape, 1, 1, true;
132            DynamicTruncate, "dynamic_truncate", Dynamic, SameAny, 2, 2, true;
133            PadToMatch, "pad_to_match", Dynamic, SameAny, 2, 2, true;
134            Constant, "constant", Host, Constant, 0, 0, true;
135        }
136    };
137}
138
139macro_rules! define_kind {
140    ($( $variant:ident, $name:literal, $category:ident, $policy:ident, $min:expr, $max:expr, $host:expr; )*) => {
141        /// Catalog key for a core primitive operation.
142        ///
143        /// # Examples
144        ///
145        /// ```rust
146        /// use tenferro_core_ops::{descriptor, PrimitiveOpKind};
147        ///
148        /// assert_eq!(descriptor(PrimitiveOpKind::Add).name, "add");
149        /// ```
150        #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
151        pub enum PrimitiveOpKind {
152            $( $variant, )*
153        }
154
155        impl PrimitiveOpKind {
156            /// Number of primitive operation kinds in the catalog.
157            ///
158            /// # Examples
159            ///
160            /// ```rust
161            /// use tenferro_core_ops::PrimitiveOpKind;
162            ///
163            /// assert!(PrimitiveOpKind::COUNT > 0);
164            /// ```
165            pub const COUNT: usize = [$(PrimitiveOpKind::$variant),*].len();
166
167            /// Return this kind's dense catalog index.
168            ///
169            /// # Examples
170            ///
171            /// ```rust
172            /// use tenferro_core_ops::PrimitiveOpKind;
173            ///
174            /// assert_eq!(PrimitiveOpKind::Add.as_index(), 0);
175            /// ```
176            pub const fn as_index(self) -> usize {
177                self as usize
178            }
179        }
180    };
181}
182
183primitive_ops!(define_kind);
184
185macro_rules! define_descriptors {
186    ($( $variant:ident, $name:literal, $category:ident, $policy:ident, $min:expr, $max:expr, $host:expr; )*) => {
187        const DESCRIPTORS: &[PrimitiveOpDescriptor] = &[
188            $(
189                PrimitiveOpDescriptor {
190                    kind: PrimitiveOpKind::$variant,
191                    name: $name,
192                    category: OpCategory::$category,
193                    dtype_policy: DTypePolicy::$policy,
194                    min_inputs: $min,
195                    max_inputs: $max,
196                    host_only: $host,
197                },
198            )*
199        ];
200
201        /// Return the descriptor for a primitive operation kind.
202        ///
203        /// # Examples
204        ///
205        /// ```rust
206        /// use tenferro_core_ops::{descriptor, PrimitiveOpKind};
207        ///
208        /// assert_eq!(descriptor(PrimitiveOpKind::Add).name, "add");
209        /// ```
210        pub fn descriptor(kind: PrimitiveOpKind) -> &'static PrimitiveOpDescriptor {
211            match kind {
212                $(
213                    PrimitiveOpKind::$variant => &DESCRIPTORS[PrimitiveOpKind::$variant as usize],
214                )*
215            }
216        }
217    };
218}
219
220primitive_ops!(define_descriptors);
221
222/// Return all core primitive operation descriptors in catalog order.
223///
224/// # Examples
225///
226/// ```rust
227/// use tenferro_core_ops::all_primitive_descriptors;
228///
229/// assert!(all_primitive_descriptors()
230///     .iter()
231///     .any(|descriptor| descriptor.name == "add"));
232/// ```
233pub fn all_primitive_descriptors() -> &'static [PrimitiveOpDescriptor] {
234    DESCRIPTORS
235}
236
237#[doc(hidden)]
238#[macro_export]
239macro_rules! define_std_tensor_op {
240    () => {
241        #[derive(Clone, Debug)]
242        pub enum StdTensorOp {
243            // Semiring arithmetic core
244            Add,
245            Sub,
246            Mul,
247            Neg,
248            Conj,
249            DotGeneral {
250                config: DotGeneralConfig,
251            },
252            Transpose {
253                perm: Vec<usize>,
254            },
255            Reshape {
256                to_shape: Vec<DimExpr>,
257            },
258            BroadcastInDim {
259                shape: Vec<DimExpr>,
260                dims: Vec<usize>,
261            },
262            Convert {
263                from: DType,
264                to: DType,
265            },
266            Constant {
267                dtype: DType,
268                bytes: Vec<u8>,
269            },
270            ReduceSum {
271                axes: Vec<usize>,
272            },
273            ReduceSumSquares {
274                axes: Vec<usize>,
275            },
276
277            // Elementwise (non-semiring)
278            Div,
279            Rem,
280            Abs,
281            Sign,
282            Maximum,
283            Minimum,
284            Compare(CompareDir),
285            Select,
286            Clamp,
287
288            // Analytic
289            Exp,
290            Log,
291            Sin,
292            Cos,
293            Tanh,
294            Sqrt,
295            Rsqrt,
296            Pow,
297            Expm1,
298            Log1p,
299
300            // Diagonal extraction / embedding (AD-closed pair)
301            ExtractDiag {
302                axis_a: usize,
303                axis_b: usize,
304            },
305            EmbedDiag {
306                axis_a: usize,
307                axis_b: usize,
308            },
309            Tril {
310                k: i64,
311            },
312            Triu {
313                k: i64,
314            },
315
316            // Indexing
317            Gather(GatherConfig),
318            GatherDynamicSliceSizes {
319                offset_dims: Vec<usize>,
320                collapsed_slice_dims: Vec<usize>,
321                start_index_map: Vec<usize>,
322                index_vector_dim: usize,
323                slice_sizes: Vec<DimExpr>,
324            },
325            Scatter(ScatterConfig),
326            Slice(SliceConfig),
327            DynamicSlice {
328                slice_sizes: Vec<usize>,
329            },
330            DynamicUpdateSlice,
331            Pad(PadConfig),
332            Concatenate {
333                axis: usize,
334                input_count: usize,
335            },
336            Reverse {
337                axes: Vec<usize>,
338            },
339            ShapeOf {
340                axis: usize,
341            },
342            DynamicTruncate {
343                axis: usize,
344            },
345            PadToMatch {
346                axis: usize,
347            },
348
349            // Reductions
350            ReduceProd {
351                axes: Vec<usize>,
352            },
353            ReduceMax {
354                axes: Vec<usize>,
355            },
356            ReduceMin {
357                axes: Vec<usize>,
358            },
359
360            /// Out-of-tree extension carrier.
361            ///
362            /// See [`crate::ext_op`] and `docs/spec/extension-op.md`. Identity,
363            /// hashing, equality, arity, shape inference, and AD rules are delegated
364            /// to the inner [`ExtensionOp`] trait object.
365            Extension(Arc<dyn ExtensionOp>),
366        }
367
368        impl StdTensorOp {
369            /// Return the core primitive catalog kind for this graph operation.
370            ///
371            /// Extension operations do not claim a core primitive kind; they are
372            /// dispatched through their extension family id instead.
373            ///
374            /// # Examples
375            ///
376            /// ```rust
377            /// use tenferro_core_ops::PrimitiveOpKind;
378            /// use tenferro_ops::std_tensor_op::StdTensorOp;
379            ///
380            /// assert_eq!(StdTensorOp::Add.primitive_kind(), Some(PrimitiveOpKind::Add));
381            /// ```
382            pub fn primitive_kind(&self) -> Option<$crate::PrimitiveOpKind> {
383                let kind = match self {
384                    Self::Add => $crate::PrimitiveOpKind::Add,
385                    Self::Sub => $crate::PrimitiveOpKind::Sub,
386                    Self::Mul => $crate::PrimitiveOpKind::Mul,
387                    Self::Neg => $crate::PrimitiveOpKind::Neg,
388                    Self::Conj => $crate::PrimitiveOpKind::Conj,
389                    Self::DotGeneral { .. } => $crate::PrimitiveOpKind::DotGeneral,
390                    Self::Transpose { .. } => $crate::PrimitiveOpKind::Transpose,
391                    Self::Reshape { .. } => $crate::PrimitiveOpKind::Reshape,
392                    Self::BroadcastInDim { .. } => $crate::PrimitiveOpKind::BroadcastInDim,
393                    Self::Convert { .. } => $crate::PrimitiveOpKind::Convert,
394                    Self::Constant { .. } => $crate::PrimitiveOpKind::Constant,
395                    Self::ReduceSum { .. } => $crate::PrimitiveOpKind::ReduceSum,
396                    Self::ReduceSumSquares { .. } => $crate::PrimitiveOpKind::ReduceSumSquares,
397                    Self::Div => $crate::PrimitiveOpKind::Div,
398                    Self::Rem => $crate::PrimitiveOpKind::Rem,
399                    Self::Abs => $crate::PrimitiveOpKind::Abs,
400                    Self::Sign => $crate::PrimitiveOpKind::Sign,
401                    Self::Maximum => $crate::PrimitiveOpKind::Maximum,
402                    Self::Minimum => $crate::PrimitiveOpKind::Minimum,
403                    Self::Compare(_) => $crate::PrimitiveOpKind::Compare,
404                    Self::Select => $crate::PrimitiveOpKind::Select,
405                    Self::Clamp => $crate::PrimitiveOpKind::Clamp,
406                    Self::Exp => $crate::PrimitiveOpKind::Exp,
407                    Self::Log => $crate::PrimitiveOpKind::Log,
408                    Self::Sin => $crate::PrimitiveOpKind::Sin,
409                    Self::Cos => $crate::PrimitiveOpKind::Cos,
410                    Self::Tanh => $crate::PrimitiveOpKind::Tanh,
411                    Self::Sqrt => $crate::PrimitiveOpKind::Sqrt,
412                    Self::Rsqrt => $crate::PrimitiveOpKind::Rsqrt,
413                    Self::Pow => $crate::PrimitiveOpKind::Pow,
414                    Self::Expm1 => $crate::PrimitiveOpKind::Expm1,
415                    Self::Log1p => $crate::PrimitiveOpKind::Log1p,
416                    Self::ExtractDiag { .. } => $crate::PrimitiveOpKind::ExtractDiag,
417                    Self::EmbedDiag { .. } => $crate::PrimitiveOpKind::EmbedDiag,
418                    Self::Tril { .. } => $crate::PrimitiveOpKind::Tril,
419                    Self::Triu { .. } => $crate::PrimitiveOpKind::Triu,
420                    Self::Gather(_) => $crate::PrimitiveOpKind::Gather,
421                    Self::GatherDynamicSliceSizes { .. } => {
422                        $crate::PrimitiveOpKind::GatherDynamicSliceSizes
423                    }
424                    Self::Scatter(_) => $crate::PrimitiveOpKind::Scatter,
425                    Self::Slice(_) => $crate::PrimitiveOpKind::Slice,
426                    Self::DynamicSlice { .. } => $crate::PrimitiveOpKind::DynamicSlice,
427                    Self::DynamicUpdateSlice => $crate::PrimitiveOpKind::DynamicUpdateSlice,
428                    Self::Pad(_) => $crate::PrimitiveOpKind::Pad,
429                    Self::Concatenate { .. } => $crate::PrimitiveOpKind::Concatenate,
430                    Self::Reverse { .. } => $crate::PrimitiveOpKind::Reverse,
431                    Self::ShapeOf { .. } => $crate::PrimitiveOpKind::ShapeOf,
432                    Self::DynamicTruncate { .. } => $crate::PrimitiveOpKind::DynamicTruncate,
433                    Self::PadToMatch { .. } => $crate::PrimitiveOpKind::PadToMatch,
434                    Self::ReduceProd { .. } => $crate::PrimitiveOpKind::ReduceProd,
435                    Self::ReduceMax { .. } => $crate::PrimitiveOpKind::ReduceMax,
436                    Self::ReduceMin { .. } => $crate::PrimitiveOpKind::ReduceMin,
437                    Self::Extension(_) => return None,
438                };
439                Some(kind)
440            }
441
442            #[cfg(test)]
443            pub(crate) fn sample_from_kind(kind: $crate::PrimitiveOpKind) -> Self {
444                match kind {
445                    $crate::PrimitiveOpKind::Add => Self::Add,
446                    $crate::PrimitiveOpKind::Sub => Self::Sub,
447                    $crate::PrimitiveOpKind::Mul => Self::Mul,
448                    $crate::PrimitiveOpKind::Neg => Self::Neg,
449                    $crate::PrimitiveOpKind::Conj => Self::Conj,
450                    $crate::PrimitiveOpKind::DotGeneral => Self::DotGeneral {
451                        config: DotGeneralConfig {
452                            lhs_contracting_dims: vec![0],
453                            rhs_contracting_dims: vec![0],
454                            lhs_batch_dims: vec![],
455                            rhs_batch_dims: vec![],
456                        },
457                    },
458                    $crate::PrimitiveOpKind::Transpose => Self::Transpose { perm: vec![0] },
459                    $crate::PrimitiveOpKind::Reshape => Self::Reshape {
460                        to_shape: vec![DimExpr::Const(1)],
461                    },
462                    $crate::PrimitiveOpKind::BroadcastInDim => Self::BroadcastInDim {
463                        shape: vec![DimExpr::Const(1)],
464                        dims: vec![0],
465                    },
466                    $crate::PrimitiveOpKind::Convert => Self::Convert {
467                        from: DType::F32,
468                        to: DType::F64,
469                    },
470                    $crate::PrimitiveOpKind::Constant => Self::Constant {
471                        dtype: DType::F64,
472                        bytes: 0.0_f64.to_le_bytes().to_vec(),
473                    },
474                    $crate::PrimitiveOpKind::ReduceSum => Self::ReduceSum { axes: vec![0] },
475                    $crate::PrimitiveOpKind::ReduceSumSquares => {
476                        Self::ReduceSumSquares { axes: vec![0] }
477                    }
478                    $crate::PrimitiveOpKind::Div => Self::Div,
479                    $crate::PrimitiveOpKind::Rem => Self::Rem,
480                    $crate::PrimitiveOpKind::Abs => Self::Abs,
481                    $crate::PrimitiveOpKind::Sign => Self::Sign,
482                    $crate::PrimitiveOpKind::Maximum => Self::Maximum,
483                    $crate::PrimitiveOpKind::Minimum => Self::Minimum,
484                    $crate::PrimitiveOpKind::Compare => Self::Compare(CompareDir::Eq),
485                    $crate::PrimitiveOpKind::Select => Self::Select,
486                    $crate::PrimitiveOpKind::Clamp => Self::Clamp,
487                    $crate::PrimitiveOpKind::Exp => Self::Exp,
488                    $crate::PrimitiveOpKind::Log => Self::Log,
489                    $crate::PrimitiveOpKind::Sin => Self::Sin,
490                    $crate::PrimitiveOpKind::Cos => Self::Cos,
491                    $crate::PrimitiveOpKind::Tanh => Self::Tanh,
492                    $crate::PrimitiveOpKind::Sqrt => Self::Sqrt,
493                    $crate::PrimitiveOpKind::Rsqrt => Self::Rsqrt,
494                    $crate::PrimitiveOpKind::Pow => Self::Pow,
495                    $crate::PrimitiveOpKind::Expm1 => Self::Expm1,
496                    $crate::PrimitiveOpKind::Log1p => Self::Log1p,
497                    $crate::PrimitiveOpKind::ExtractDiag => Self::ExtractDiag {
498                        axis_a: 0,
499                        axis_b: 1,
500                    },
501                    $crate::PrimitiveOpKind::EmbedDiag => Self::EmbedDiag {
502                        axis_a: 0,
503                        axis_b: 1,
504                    },
505                    $crate::PrimitiveOpKind::Tril => Self::Tril { k: 0 },
506                    $crate::PrimitiveOpKind::Triu => Self::Triu { k: 0 },
507                    $crate::PrimitiveOpKind::Gather => Self::Gather(GatherConfig {
508                        offset_dims: vec![],
509                        collapsed_slice_dims: vec![0],
510                        start_index_map: vec![0],
511                        index_vector_dim: 1,
512                        slice_sizes: vec![1],
513                    }),
514                    $crate::PrimitiveOpKind::GatherDynamicSliceSizes => {
515                        Self::GatherDynamicSliceSizes {
516                            offset_dims: vec![],
517                            collapsed_slice_dims: vec![0],
518                            start_index_map: vec![0],
519                            index_vector_dim: 1,
520                            slice_sizes: vec![DimExpr::Const(1)],
521                        }
522                    }
523                    $crate::PrimitiveOpKind::Scatter => Self::Scatter(ScatterConfig {
524                        update_window_dims: vec![],
525                        inserted_window_dims: vec![0],
526                        scatter_dims_to_operand_dims: vec![0],
527                        index_vector_dim: 1,
528                    }),
529                    $crate::PrimitiveOpKind::Slice => Self::Slice(SliceConfig {
530                        starts: vec![0],
531                        limits: vec![1],
532                        strides: vec![1],
533                    }),
534                    $crate::PrimitiveOpKind::DynamicSlice => Self::DynamicSlice {
535                        slice_sizes: vec![1],
536                    },
537                    $crate::PrimitiveOpKind::DynamicUpdateSlice => Self::DynamicUpdateSlice,
538                    $crate::PrimitiveOpKind::Pad => Self::Pad(PadConfig {
539                        edge_padding_low: vec![0],
540                        edge_padding_high: vec![0],
541                        interior_padding: vec![0],
542                    }),
543                    $crate::PrimitiveOpKind::Concatenate => Self::Concatenate {
544                        axis: 0,
545                        input_count: 1,
546                    },
547                    $crate::PrimitiveOpKind::Reverse => Self::Reverse { axes: vec![0] },
548                    $crate::PrimitiveOpKind::ShapeOf => Self::ShapeOf { axis: 0 },
549                    $crate::PrimitiveOpKind::DynamicTruncate => Self::DynamicTruncate { axis: 0 },
550                    $crate::PrimitiveOpKind::PadToMatch => Self::PadToMatch { axis: 0 },
551                    $crate::PrimitiveOpKind::ReduceProd => Self::ReduceProd { axes: vec![0] },
552                    $crate::PrimitiveOpKind::ReduceMax => Self::ReduceMax { axes: vec![0] },
553                    $crate::PrimitiveOpKind::ReduceMin => Self::ReduceMin { axes: vec![0] },
554                }
555            }
556        }
557    };
558}
559
560#[doc(hidden)]
561#[macro_export]
562macro_rules! define_elementwise_fusion_op {
563    () => {
564        /// Elementwise op kinds supported by backend fusion implementations.
565        #[doc(hidden)]
566        #[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
567        pub enum ElementwiseFusionOp {
568            Add,
569            Multiply,
570            Negate,
571            Conj,
572            Divide,
573            Remainder,
574            Abs,
575            Maximum,
576            Minimum,
577            Clamp,
578            Exp,
579            Log,
580            Sin,
581            Cos,
582            Tanh,
583            Sqrt,
584            Rsqrt,
585            Pow,
586            Expm1,
587            Log1p,
588        }
589
590        #[cfg(test)]
591        impl ElementwiseFusionOp {
592            pub(crate) fn iter() -> impl Iterator<Item = Self> {
593                [
594                    Self::Add,
595                    Self::Multiply,
596                    Self::Negate,
597                    Self::Conj,
598                    Self::Divide,
599                    Self::Remainder,
600                    Self::Abs,
601                    Self::Maximum,
602                    Self::Minimum,
603                    Self::Clamp,
604                    Self::Exp,
605                    Self::Log,
606                    Self::Sin,
607                    Self::Cos,
608                    Self::Tanh,
609                    Self::Sqrt,
610                    Self::Rsqrt,
611                    Self::Pow,
612                    Self::Expm1,
613                    Self::Log1p,
614                ]
615                .into_iter()
616            }
617
618            pub(crate) fn from_primitive_kind(kind: $crate::PrimitiveOpKind) -> Option<Self> {
619                match kind {
620                    $crate::PrimitiveOpKind::Add => Some(Self::Add),
621                    $crate::PrimitiveOpKind::Mul => Some(Self::Multiply),
622                    $crate::PrimitiveOpKind::Neg => Some(Self::Negate),
623                    $crate::PrimitiveOpKind::Conj => Some(Self::Conj),
624                    $crate::PrimitiveOpKind::Div => Some(Self::Divide),
625                    $crate::PrimitiveOpKind::Rem => Some(Self::Remainder),
626                    $crate::PrimitiveOpKind::Abs => Some(Self::Abs),
627                    $crate::PrimitiveOpKind::Maximum => Some(Self::Maximum),
628                    $crate::PrimitiveOpKind::Minimum => Some(Self::Minimum),
629                    $crate::PrimitiveOpKind::Clamp => Some(Self::Clamp),
630                    $crate::PrimitiveOpKind::Exp => Some(Self::Exp),
631                    $crate::PrimitiveOpKind::Log => Some(Self::Log),
632                    $crate::PrimitiveOpKind::Sin => Some(Self::Sin),
633                    $crate::PrimitiveOpKind::Cos => Some(Self::Cos),
634                    $crate::PrimitiveOpKind::Tanh => Some(Self::Tanh),
635                    $crate::PrimitiveOpKind::Sqrt => Some(Self::Sqrt),
636                    $crate::PrimitiveOpKind::Rsqrt => Some(Self::Rsqrt),
637                    $crate::PrimitiveOpKind::Pow => Some(Self::Pow),
638                    $crate::PrimitiveOpKind::Expm1 => Some(Self::Expm1),
639                    $crate::PrimitiveOpKind::Log1p => Some(Self::Log1p),
640                    _ => None,
641                }
642            }
643
644            pub(crate) fn primitive_kind(self) -> $crate::PrimitiveOpKind {
645                match self {
646                    Self::Add => $crate::PrimitiveOpKind::Add,
647                    Self::Multiply => $crate::PrimitiveOpKind::Mul,
648                    Self::Negate => $crate::PrimitiveOpKind::Neg,
649                    Self::Conj => $crate::PrimitiveOpKind::Conj,
650                    Self::Divide => $crate::PrimitiveOpKind::Div,
651                    Self::Remainder => $crate::PrimitiveOpKind::Rem,
652                    Self::Abs => $crate::PrimitiveOpKind::Abs,
653                    Self::Maximum => $crate::PrimitiveOpKind::Maximum,
654                    Self::Minimum => $crate::PrimitiveOpKind::Minimum,
655                    Self::Clamp => $crate::PrimitiveOpKind::Clamp,
656                    Self::Exp => $crate::PrimitiveOpKind::Exp,
657                    Self::Log => $crate::PrimitiveOpKind::Log,
658                    Self::Sin => $crate::PrimitiveOpKind::Sin,
659                    Self::Cos => $crate::PrimitiveOpKind::Cos,
660                    Self::Tanh => $crate::PrimitiveOpKind::Tanh,
661                    Self::Sqrt => $crate::PrimitiveOpKind::Sqrt,
662                    Self::Rsqrt => $crate::PrimitiveOpKind::Rsqrt,
663                    Self::Pow => $crate::PrimitiveOpKind::Pow,
664                    Self::Expm1 => $crate::PrimitiveOpKind::Expm1,
665                    Self::Log1p => $crate::PrimitiveOpKind::Log1p,
666                }
667            }
668        }
669    };
670}
671
672#[doc(hidden)]
673#[macro_export]
674macro_rules! define_exec_op {
675    () => {
676        #[derive(Clone, Debug)]
677        pub enum ExecOp {
678            Transpose {
679                perm: Vec<usize>,
680            },
681            Reshape {
682                shape: Vec<DimExpr>,
683            },
684            BroadcastInDim {
685                shape: Vec<DimExpr>,
686                dims: Vec<usize>,
687            },
688            Convert {
689                to: DType,
690            },
691            Constant {
692                dtype: DType,
693                bytes: Vec<u8>,
694            },
695            DotGeneral(DotGeneralConfig),
696            DotGeneralWithConj {
697                config: DotGeneralConfig,
698                lhs_conj: bool,
699                rhs_conj: bool,
700            },
701            ReduceSum {
702                axes: Vec<usize>,
703            },
704            ReduceSumSquares {
705                axes: Vec<usize>,
706            },
707            ExtractDiag {
708                axis_a: usize,
709                axis_b: usize,
710            },
711            EmbedDiag {
712                axis_a: usize,
713                axis_b: usize,
714            },
715            Tril {
716                k: i64,
717            },
718            Triu {
719                k: i64,
720            },
721            Add,
722            Subtract,
723            Multiply,
724            Negate,
725            Conj,
726            Divide,
727            Remainder,
728            Abs,
729            Sign,
730            Maximum,
731            Minimum,
732            Compare(CompareDir),
733            Select,
734            Clamp,
735            Exp,
736            Log,
737            Sin,
738            Cos,
739            Tanh,
740            Sqrt,
741            Rsqrt,
742            Pow,
743            Expm1,
744            Log1p,
745            Gather(GatherConfig),
746            GatherDynamicSliceSizes {
747                offset_dims: Vec<usize>,
748                collapsed_slice_dims: Vec<usize>,
749                start_index_map: Vec<usize>,
750                index_vector_dim: usize,
751                slice_sizes: Vec<DimExpr>,
752            },
753            Scatter(ScatterConfig),
754            Slice(SliceConfig),
755            DynamicSlice {
756                slice_sizes: Vec<usize>,
757            },
758            DynamicUpdateSlice,
759            Pad(PadConfig),
760            Concatenate {
761                axis: usize,
762            },
763            Reverse {
764                axes: Vec<usize>,
765            },
766            ShapeOf {
767                axis: usize,
768            },
769            DynamicTruncate {
770                axis: usize,
771            },
772            PadToMatch {
773                axis: usize,
774            },
775            ReduceProd {
776                axes: Vec<usize>,
777            },
778            ReduceMax {
779                axes: Vec<usize>,
780            },
781            ReduceMin {
782                axes: Vec<usize>,
783            },
784            /// Out-of-tree extension carrier in the execution IR.
785            ///
786            /// Payload and dispatch are defined by the inner [`ExtensionOp`]. The
787            /// execution pipeline treats extensions as single-instruction FFI
788            /// boundaries (spec Section 8): no elementwise fusion, and dispatch is
789            /// routed through the executor's registered extension runtime.
790            Extension(Arc<dyn ExtensionOp>),
791        }
792
793        impl ExecOp {
794            pub(crate) fn primitive_kind(&self) -> Option<$crate::PrimitiveOpKind> {
795                let kind = match self {
796                    Self::Transpose { .. } => $crate::PrimitiveOpKind::Transpose,
797                    Self::Reshape { .. } => $crate::PrimitiveOpKind::Reshape,
798                    Self::BroadcastInDim { .. } => $crate::PrimitiveOpKind::BroadcastInDim,
799                    Self::Convert { .. } => $crate::PrimitiveOpKind::Convert,
800                    Self::Constant { .. } => $crate::PrimitiveOpKind::Constant,
801                    Self::DotGeneral(_) | Self::DotGeneralWithConj { .. } => {
802                        $crate::PrimitiveOpKind::DotGeneral
803                    }
804                    Self::ReduceSum { .. } => $crate::PrimitiveOpKind::ReduceSum,
805                    Self::ReduceSumSquares { .. } => $crate::PrimitiveOpKind::ReduceSumSquares,
806                    Self::ExtractDiag { .. } => $crate::PrimitiveOpKind::ExtractDiag,
807                    Self::EmbedDiag { .. } => $crate::PrimitiveOpKind::EmbedDiag,
808                    Self::Tril { .. } => $crate::PrimitiveOpKind::Tril,
809                    Self::Triu { .. } => $crate::PrimitiveOpKind::Triu,
810                    Self::Add => $crate::PrimitiveOpKind::Add,
811                    Self::Subtract => $crate::PrimitiveOpKind::Sub,
812                    Self::Multiply => $crate::PrimitiveOpKind::Mul,
813                    Self::Negate => $crate::PrimitiveOpKind::Neg,
814                    Self::Conj => $crate::PrimitiveOpKind::Conj,
815                    Self::Divide => $crate::PrimitiveOpKind::Div,
816                    Self::Remainder => $crate::PrimitiveOpKind::Rem,
817                    Self::Abs => $crate::PrimitiveOpKind::Abs,
818                    Self::Sign => $crate::PrimitiveOpKind::Sign,
819                    Self::Maximum => $crate::PrimitiveOpKind::Maximum,
820                    Self::Minimum => $crate::PrimitiveOpKind::Minimum,
821                    Self::Compare(_) => $crate::PrimitiveOpKind::Compare,
822                    Self::Select => $crate::PrimitiveOpKind::Select,
823                    Self::Clamp => $crate::PrimitiveOpKind::Clamp,
824                    Self::Exp => $crate::PrimitiveOpKind::Exp,
825                    Self::Log => $crate::PrimitiveOpKind::Log,
826                    Self::Sin => $crate::PrimitiveOpKind::Sin,
827                    Self::Cos => $crate::PrimitiveOpKind::Cos,
828                    Self::Tanh => $crate::PrimitiveOpKind::Tanh,
829                    Self::Sqrt => $crate::PrimitiveOpKind::Sqrt,
830                    Self::Rsqrt => $crate::PrimitiveOpKind::Rsqrt,
831                    Self::Pow => $crate::PrimitiveOpKind::Pow,
832                    Self::Expm1 => $crate::PrimitiveOpKind::Expm1,
833                    Self::Log1p => $crate::PrimitiveOpKind::Log1p,
834                    Self::Gather(_) => $crate::PrimitiveOpKind::Gather,
835                    Self::GatherDynamicSliceSizes { .. } => {
836                        $crate::PrimitiveOpKind::GatherDynamicSliceSizes
837                    }
838                    Self::Scatter(_) => $crate::PrimitiveOpKind::Scatter,
839                    Self::Slice(_) => $crate::PrimitiveOpKind::Slice,
840                    Self::DynamicSlice { .. } => $crate::PrimitiveOpKind::DynamicSlice,
841                    Self::DynamicUpdateSlice => $crate::PrimitiveOpKind::DynamicUpdateSlice,
842                    Self::Pad(_) => $crate::PrimitiveOpKind::Pad,
843                    Self::Concatenate { .. } => $crate::PrimitiveOpKind::Concatenate,
844                    Self::Reverse { .. } => $crate::PrimitiveOpKind::Reverse,
845                    Self::ShapeOf { .. } => $crate::PrimitiveOpKind::ShapeOf,
846                    Self::DynamicTruncate { .. } => $crate::PrimitiveOpKind::DynamicTruncate,
847                    Self::PadToMatch { .. } => $crate::PrimitiveOpKind::PadToMatch,
848                    Self::ReduceProd { .. } => $crate::PrimitiveOpKind::ReduceProd,
849                    Self::ReduceMax { .. } => $crate::PrimitiveOpKind::ReduceMax,
850                    Self::ReduceMin { .. } => $crate::PrimitiveOpKind::ReduceMin,
851                    Self::Extension(_) => return None,
852                };
853                Some(kind)
854            }
855
856            pub(crate) fn from_std_tensor_op(
857                op: &tenferro_ops::std_tensor_op::StdTensorOp,
858            ) -> Self {
859                match op {
860                    tenferro_ops::std_tensor_op::StdTensorOp::Add => Self::Add,
861                    tenferro_ops::std_tensor_op::StdTensorOp::Sub => Self::Subtract,
862                    tenferro_ops::std_tensor_op::StdTensorOp::Mul => Self::Multiply,
863                    tenferro_ops::std_tensor_op::StdTensorOp::Neg => Self::Negate,
864                    tenferro_ops::std_tensor_op::StdTensorOp::Conj => Self::Conj,
865                    tenferro_ops::std_tensor_op::StdTensorOp::Div => Self::Divide,
866                    tenferro_ops::std_tensor_op::StdTensorOp::Rem => Self::Remainder,
867                    tenferro_ops::std_tensor_op::StdTensorOp::Abs => Self::Abs,
868                    tenferro_ops::std_tensor_op::StdTensorOp::Sign => Self::Sign,
869                    tenferro_ops::std_tensor_op::StdTensorOp::Maximum => Self::Maximum,
870                    tenferro_ops::std_tensor_op::StdTensorOp::Minimum => Self::Minimum,
871                    tenferro_ops::std_tensor_op::StdTensorOp::Compare(dir) => {
872                        Self::Compare(dir.clone())
873                    }
874                    tenferro_ops::std_tensor_op::StdTensorOp::Select => Self::Select,
875                    tenferro_ops::std_tensor_op::StdTensorOp::Clamp => Self::Clamp,
876                    tenferro_ops::std_tensor_op::StdTensorOp::Exp => Self::Exp,
877                    tenferro_ops::std_tensor_op::StdTensorOp::Log => Self::Log,
878                    tenferro_ops::std_tensor_op::StdTensorOp::Sin => Self::Sin,
879                    tenferro_ops::std_tensor_op::StdTensorOp::Cos => Self::Cos,
880                    tenferro_ops::std_tensor_op::StdTensorOp::Tanh => Self::Tanh,
881                    tenferro_ops::std_tensor_op::StdTensorOp::Sqrt => Self::Sqrt,
882                    tenferro_ops::std_tensor_op::StdTensorOp::Rsqrt => Self::Rsqrt,
883                    tenferro_ops::std_tensor_op::StdTensorOp::Pow => Self::Pow,
884                    tenferro_ops::std_tensor_op::StdTensorOp::Expm1 => Self::Expm1,
885                    tenferro_ops::std_tensor_op::StdTensorOp::Log1p => Self::Log1p,
886                    tenferro_ops::std_tensor_op::StdTensorOp::Transpose { perm } => {
887                        Self::Transpose { perm: perm.clone() }
888                    }
889                    tenferro_ops::std_tensor_op::StdTensorOp::Reshape { to_shape } => {
890                        Self::Reshape {
891                            shape: to_shape.clone(),
892                        }
893                    }
894                    tenferro_ops::std_tensor_op::StdTensorOp::BroadcastInDim { shape, dims } => {
895                        Self::BroadcastInDim {
896                            shape: shape.clone(),
897                            dims: dims.clone(),
898                        }
899                    }
900                    tenferro_ops::std_tensor_op::StdTensorOp::Convert { to, .. } => {
901                        Self::Convert { to: *to }
902                    }
903                    tenferro_ops::std_tensor_op::StdTensorOp::Constant { dtype, bytes } => {
904                        Self::Constant {
905                            dtype: *dtype,
906                            bytes: bytes.clone(),
907                        }
908                    }
909                    tenferro_ops::std_tensor_op::StdTensorOp::DotGeneral { config } => {
910                        Self::DotGeneral(config.clone())
911                    }
912                    tenferro_ops::std_tensor_op::StdTensorOp::ReduceSum { axes } => {
913                        Self::ReduceSum { axes: axes.clone() }
914                    }
915                    tenferro_ops::std_tensor_op::StdTensorOp::ReduceSumSquares { axes } => {
916                        Self::ReduceSumSquares { axes: axes.clone() }
917                    }
918                    tenferro_ops::std_tensor_op::StdTensorOp::ReduceProd { axes } => {
919                        Self::ReduceProd { axes: axes.clone() }
920                    }
921                    tenferro_ops::std_tensor_op::StdTensorOp::ReduceMax { axes } => {
922                        Self::ReduceMax { axes: axes.clone() }
923                    }
924                    tenferro_ops::std_tensor_op::StdTensorOp::ReduceMin { axes } => {
925                        Self::ReduceMin { axes: axes.clone() }
926                    }
927                    tenferro_ops::std_tensor_op::StdTensorOp::ExtractDiag { axis_a, axis_b } => {
928                        Self::ExtractDiag {
929                            axis_a: *axis_a,
930                            axis_b: *axis_b,
931                        }
932                    }
933                    tenferro_ops::std_tensor_op::StdTensorOp::EmbedDiag { axis_a, axis_b } => {
934                        Self::EmbedDiag {
935                            axis_a: *axis_a,
936                            axis_b: *axis_b,
937                        }
938                    }
939                    tenferro_ops::std_tensor_op::StdTensorOp::Tril { k } => Self::Tril { k: *k },
940                    tenferro_ops::std_tensor_op::StdTensorOp::Triu { k } => Self::Triu { k: *k },
941                    tenferro_ops::std_tensor_op::StdTensorOp::Gather(config) => {
942                        Self::Gather(config.clone())
943                    }
944                    tenferro_ops::std_tensor_op::StdTensorOp::GatherDynamicSliceSizes {
945                        offset_dims,
946                        collapsed_slice_dims,
947                        start_index_map,
948                        index_vector_dim,
949                        slice_sizes,
950                    } => Self::GatherDynamicSliceSizes {
951                        offset_dims: offset_dims.clone(),
952                        collapsed_slice_dims: collapsed_slice_dims.clone(),
953                        start_index_map: start_index_map.clone(),
954                        index_vector_dim: *index_vector_dim,
955                        slice_sizes: slice_sizes.clone(),
956                    },
957                    tenferro_ops::std_tensor_op::StdTensorOp::Scatter(config) => {
958                        Self::Scatter(config.clone())
959                    }
960                    tenferro_ops::std_tensor_op::StdTensorOp::Slice(config) => {
961                        Self::Slice(config.clone())
962                    }
963                    tenferro_ops::std_tensor_op::StdTensorOp::DynamicSlice { slice_sizes } => {
964                        Self::DynamicSlice {
965                            slice_sizes: slice_sizes.clone(),
966                        }
967                    }
968                    tenferro_ops::std_tensor_op::StdTensorOp::DynamicUpdateSlice => {
969                        Self::DynamicUpdateSlice
970                    }
971                    tenferro_ops::std_tensor_op::StdTensorOp::Pad(config) => {
972                        Self::Pad(config.clone())
973                    }
974                    tenferro_ops::std_tensor_op::StdTensorOp::Concatenate { axis, .. } => {
975                        Self::Concatenate { axis: *axis }
976                    }
977                    tenferro_ops::std_tensor_op::StdTensorOp::Reverse { axes } => {
978                        Self::Reverse { axes: axes.clone() }
979                    }
980                    tenferro_ops::std_tensor_op::StdTensorOp::ShapeOf { axis } => {
981                        Self::ShapeOf { axis: *axis }
982                    }
983                    tenferro_ops::std_tensor_op::StdTensorOp::DynamicTruncate { axis } => {
984                        Self::DynamicTruncate { axis: *axis }
985                    }
986                    tenferro_ops::std_tensor_op::StdTensorOp::PadToMatch { axis } => {
987                        Self::PadToMatch { axis: *axis }
988                    }
989                    tenferro_ops::std_tensor_op::StdTensorOp::Extension(op) => {
990                        Self::Extension(op.clone())
991                    }
992                }
993            }
994
995            pub(crate) fn elementwise_fusion_op(&self) -> Option<ElementwiseFusionOp> {
996                match self {
997                    Self::Add => Some(ElementwiseFusionOp::Add),
998                    Self::Multiply => Some(ElementwiseFusionOp::Multiply),
999                    Self::Negate => Some(ElementwiseFusionOp::Negate),
1000                    Self::Conj => Some(ElementwiseFusionOp::Conj),
1001                    Self::Divide => Some(ElementwiseFusionOp::Divide),
1002                    Self::Abs => Some(ElementwiseFusionOp::Abs),
1003                    Self::Maximum => Some(ElementwiseFusionOp::Maximum),
1004                    Self::Minimum => Some(ElementwiseFusionOp::Minimum),
1005                    Self::Clamp => Some(ElementwiseFusionOp::Clamp),
1006                    Self::Exp => Some(ElementwiseFusionOp::Exp),
1007                    Self::Log => Some(ElementwiseFusionOp::Log),
1008                    Self::Sin => Some(ElementwiseFusionOp::Sin),
1009                    Self::Cos => Some(ElementwiseFusionOp::Cos),
1010                    Self::Tanh => Some(ElementwiseFusionOp::Tanh),
1011                    Self::Sqrt => Some(ElementwiseFusionOp::Sqrt),
1012                    Self::Rsqrt => Some(ElementwiseFusionOp::Rsqrt),
1013                    Self::Pow => Some(ElementwiseFusionOp::Pow),
1014                    Self::Expm1 => Some(ElementwiseFusionOp::Expm1),
1015                    Self::Log1p => Some(ElementwiseFusionOp::Log1p),
1016                    _ => None,
1017                }
1018            }
1019
1020            #[cfg(test)]
1021            pub(crate) fn input_arity_bounds(&self) -> Option<(u8, u8)> {
1022                self.primitive_kind().map(|kind| {
1023                    let descriptor = $crate::descriptor(kind);
1024                    (descriptor.min_inputs, descriptor.max_inputs)
1025                })
1026            }
1027
1028            #[cfg(test)]
1029            pub(crate) fn sample_from_kind(kind: $crate::PrimitiveOpKind) -> Self {
1030                match kind {
1031                    $crate::PrimitiveOpKind::Transpose => Self::Transpose { perm: vec![0] },
1032                    $crate::PrimitiveOpKind::Reshape => Self::Reshape {
1033                        shape: vec![DimExpr::Const(1)],
1034                    },
1035                    $crate::PrimitiveOpKind::BroadcastInDim => Self::BroadcastInDim {
1036                        shape: vec![DimExpr::Const(1)],
1037                        dims: vec![0],
1038                    },
1039                    $crate::PrimitiveOpKind::Convert => Self::Convert { to: DType::F64 },
1040                    $crate::PrimitiveOpKind::Constant => Self::Constant {
1041                        dtype: DType::F64,
1042                        bytes: 0.0_f64.to_le_bytes().to_vec(),
1043                    },
1044                    $crate::PrimitiveOpKind::DotGeneral => Self::DotGeneral(DotGeneralConfig {
1045                        lhs_contracting_dims: vec![0],
1046                        rhs_contracting_dims: vec![0],
1047                        lhs_batch_dims: vec![],
1048                        rhs_batch_dims: vec![],
1049                    }),
1050                    $crate::PrimitiveOpKind::ReduceSum => Self::ReduceSum { axes: vec![0] },
1051                    $crate::PrimitiveOpKind::ReduceSumSquares => {
1052                        Self::ReduceSumSquares { axes: vec![0] }
1053                    }
1054                    $crate::PrimitiveOpKind::ExtractDiag => Self::ExtractDiag {
1055                        axis_a: 0,
1056                        axis_b: 1,
1057                    },
1058                    $crate::PrimitiveOpKind::EmbedDiag => Self::EmbedDiag {
1059                        axis_a: 0,
1060                        axis_b: 1,
1061                    },
1062                    $crate::PrimitiveOpKind::Tril => Self::Tril { k: 0 },
1063                    $crate::PrimitiveOpKind::Triu => Self::Triu { k: 0 },
1064                    $crate::PrimitiveOpKind::Add => Self::Add,
1065                    $crate::PrimitiveOpKind::Sub => Self::Subtract,
1066                    $crate::PrimitiveOpKind::Mul => Self::Multiply,
1067                    $crate::PrimitiveOpKind::Neg => Self::Negate,
1068                    $crate::PrimitiveOpKind::Conj => Self::Conj,
1069                    $crate::PrimitiveOpKind::Div => Self::Divide,
1070                    $crate::PrimitiveOpKind::Rem => Self::Remainder,
1071                    $crate::PrimitiveOpKind::Abs => Self::Abs,
1072                    $crate::PrimitiveOpKind::Sign => Self::Sign,
1073                    $crate::PrimitiveOpKind::Maximum => Self::Maximum,
1074                    $crate::PrimitiveOpKind::Minimum => Self::Minimum,
1075                    $crate::PrimitiveOpKind::Compare => Self::Compare(CompareDir::Eq),
1076                    $crate::PrimitiveOpKind::Select => Self::Select,
1077                    $crate::PrimitiveOpKind::Clamp => Self::Clamp,
1078                    $crate::PrimitiveOpKind::Exp => Self::Exp,
1079                    $crate::PrimitiveOpKind::Log => Self::Log,
1080                    $crate::PrimitiveOpKind::Sin => Self::Sin,
1081                    $crate::PrimitiveOpKind::Cos => Self::Cos,
1082                    $crate::PrimitiveOpKind::Tanh => Self::Tanh,
1083                    $crate::PrimitiveOpKind::Sqrt => Self::Sqrt,
1084                    $crate::PrimitiveOpKind::Rsqrt => Self::Rsqrt,
1085                    $crate::PrimitiveOpKind::Pow => Self::Pow,
1086                    $crate::PrimitiveOpKind::Expm1 => Self::Expm1,
1087                    $crate::PrimitiveOpKind::Log1p => Self::Log1p,
1088                    $crate::PrimitiveOpKind::Gather => Self::Gather(GatherConfig {
1089                        offset_dims: vec![],
1090                        collapsed_slice_dims: vec![0],
1091                        start_index_map: vec![0],
1092                        index_vector_dim: 1,
1093                        slice_sizes: vec![1],
1094                    }),
1095                    $crate::PrimitiveOpKind::GatherDynamicSliceSizes => {
1096                        Self::GatherDynamicSliceSizes {
1097                            offset_dims: vec![],
1098                            collapsed_slice_dims: vec![0],
1099                            start_index_map: vec![0],
1100                            index_vector_dim: 1,
1101                            slice_sizes: vec![DimExpr::Const(1)],
1102                        }
1103                    }
1104                    $crate::PrimitiveOpKind::Scatter => Self::Scatter(ScatterConfig {
1105                        update_window_dims: vec![],
1106                        inserted_window_dims: vec![0],
1107                        scatter_dims_to_operand_dims: vec![0],
1108                        index_vector_dim: 1,
1109                    }),
1110                    $crate::PrimitiveOpKind::Slice => Self::Slice(SliceConfig {
1111                        starts: vec![0],
1112                        limits: vec![1],
1113                        strides: vec![1],
1114                    }),
1115                    $crate::PrimitiveOpKind::DynamicSlice => Self::DynamicSlice {
1116                        slice_sizes: vec![1],
1117                    },
1118                    $crate::PrimitiveOpKind::DynamicUpdateSlice => Self::DynamicUpdateSlice,
1119                    $crate::PrimitiveOpKind::Pad => Self::Pad(PadConfig {
1120                        edge_padding_low: vec![0],
1121                        edge_padding_high: vec![0],
1122                        interior_padding: vec![0],
1123                    }),
1124                    $crate::PrimitiveOpKind::Concatenate => Self::Concatenate { axis: 0 },
1125                    $crate::PrimitiveOpKind::Reverse => Self::Reverse { axes: vec![0] },
1126                    $crate::PrimitiveOpKind::ShapeOf => Self::ShapeOf { axis: 0 },
1127                    $crate::PrimitiveOpKind::DynamicTruncate => Self::DynamicTruncate { axis: 0 },
1128                    $crate::PrimitiveOpKind::PadToMatch => Self::PadToMatch { axis: 0 },
1129                    $crate::PrimitiveOpKind::ReduceProd => Self::ReduceProd { axes: vec![0] },
1130                    $crate::PrimitiveOpKind::ReduceMax => Self::ReduceMax { axes: vec![0] },
1131                    $crate::PrimitiveOpKind::ReduceMin => Self::ReduceMin { axes: vec![0] },
1132                }
1133            }
1134        }
1135    };
1136}
tenferro_core_ops/catalog.rs

tenferro_core_ops/
catalog.rs
