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

tenferro_core_ops/
catalog.rs
