burn_ir/
operation.rs

1use burn_backend::ops::{AttentionModuleOptions, CausalAlignment};
2use burn_backend::tensor::IndexingUpdateOp;
3use core::hash::Hash;
4use serde::{Deserialize, Serialize};
5
6use alloc::borrow::ToOwned;
7use alloc::boxed::Box;
8use alloc::{string::String, vec::Vec};
9
10use burn_backend::{
11    DType, Distribution, Slice,
12    ops::{
13        ConvOptions, ConvTransposeOptions, CoordinateTransformMode, CubicOptions,
14        DeformConvOptions, GridSampleOptions, GridSamplePaddingMode, InterpolateMode,
15        InterpolateOptions, NearestRoundingMode, PadMode,
16    },
17    quantization::QuantScheme,
18};
19
20use crate::{ScalarIr, TensorId, TensorIr, TensorStatus};
21
22/// Visitor for mutating the components of an [`OperationIr`] in place.
23pub trait IrVisitorMut {
24    /// Visit a [`TensorIr`] mutably.
25    fn visit_tensor_mut(&mut self, _tensor: &mut TensorIr) {}
26    /// Visit a [`ScalarIr`] mutably.
27    fn visit_scalar_mut(&mut self, _scalar: &mut ScalarIr) {}
28    /// Visit a slice range mutably.
29    fn visit_range_mut(&mut self, _range: &mut Slice) {}
30}
31
32/// Custom operation in fusion stream, declaring its inputs, outputs and scalar arguments.
33#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
34pub struct CustomOpIr {
35    /// Unique identifier of the operation.
36    pub id: String,
37    /// Input tensors used in the custom operation.
38    pub inputs: Vec<TensorIr>,
39    /// Output tensors used in the custom operation.
40    pub outputs: Vec<TensorIr>,
41    /// Non-tensor scalar arguments, in declaration order.
42    ///
43    /// Carried through fusion's relativization like any other scalar (see the
44    /// `RelativeOps for CustomOpIr` impl), so a cached graph replays with fresh scalar values. The
45    /// remote backend relies on these to ship a custom op's scalar arguments to the server, where
46    /// the registered handler reads them back with [`ScalarIr::elem`].
47    pub scalars: Vec<ScalarIr>,
48}
49
50impl CustomOpIr {
51    /// Create a new custom operation intermediate representation (without scalar arguments).
52    pub fn new(id: &'static str, inputs: &[TensorIr], outputs: &[TensorIr]) -> Self {
53        Self {
54            id: id.to_owned(),
55            inputs: inputs.to_vec(),
56            outputs: outputs.to_vec(),
57            scalars: Vec::new(),
58        }
59    }
60
61    /// Create a new custom operation intermediate representation with scalar arguments.
62    pub fn with_scalars(
63        id: &'static str,
64        inputs: &[TensorIr],
65        outputs: &[TensorIr],
66        scalars: Vec<ScalarIr>,
67    ) -> Self {
68        Self {
69            id: id.to_owned(),
70            inputs: inputs.to_vec(),
71            outputs: outputs.to_vec(),
72            scalars,
73        }
74    }
75
76    /// Cast the intermediate representation, and get the in and output tensors.
77    pub fn as_fixed<const N_IN: usize, const N_OUT: usize>(
78        &self,
79    ) -> (&[TensorIr; N_IN], &[TensorIr; N_OUT]) {
80        (
81            self.inputs.as_slice().try_into().expect(
82                "Wrong number of inputs expected (expected {D}, is {}), check your implementation",
83            ),
84            self.outputs.as_slice().try_into().expect(
85                "Wrong number of outputs expected (expected {D}, is {}), check your implementation",
86            ),
87        )
88    }
89
90    fn inputs(&self) -> Box<dyn Iterator<Item = &TensorIr> + '_> {
91        Box::new(self.inputs.iter())
92    }
93
94    fn outputs(&self) -> Box<dyn Iterator<Item = &TensorIr> + '_> {
95        Box::new(self.outputs.iter())
96    }
97
98    fn visit_mut(&mut self, v: &mut impl IrVisitorMut) {
99        for t in self.inputs.iter_mut() {
100            v.visit_tensor_mut(t);
101        }
102        for t in self.outputs.iter_mut() {
103            v.visit_tensor_mut(t);
104        }
105        for s in self.scalars.iter_mut() {
106            v.visit_scalar_mut(s);
107        }
108    }
109}
110
111/// Describe all tensor operations possible.
112#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
113#[allow(clippy::large_enum_variant)]
114pub enum OperationIr {
115    /// Basic operation on a float tensor.
116    BaseFloat(BaseOperationIr),
117    /// Basic operation on an int tensor.
118    BaseInt(BaseOperationIr),
119    /// Basic operation on a bool tensor.
120    BaseBool(BaseOperationIr),
121    /// Numeric operation on a float tensor.
122    NumericFloat(DType, NumericOperationIr),
123    /// Numeric operation on an int tensor.
124    NumericInt(DType, NumericOperationIr),
125    /// Operation specific to a bool tensor.
126    Bool(BoolOperationIr),
127    /// Operation specific to an int tensor.
128    Int(IntOperationIr),
129    /// Operation specific to a float tensor.
130    Float(DType, FloatOperationIr),
131    /// Module operation.
132    Module(ModuleOperationIr),
133    /// Initialize operation.
134    Init(InitOperationIr),
135    /// A custom operation.
136    Custom(CustomOpIr),
137    /// A tensor is dropped.
138    Drop(TensorIr),
139    /// Activation function operation (relu, gelu, sigmoid, softmax, …).
140    Activation(ActivationOperationIr),
141}
142
143/// Operation intermediate representation specific to a float tensor.
144#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
145pub enum FloatOperationIr {
146    /// Operation corresponding to [exp](burn_backend::ops::FloatTensorOps::float_exp).
147    Exp(UnaryOpIr),
148    /// Operation corresponding to [log](burn_backend::ops::FloatTensorOps::float_log).
149    Log(UnaryOpIr),
150    /// Operation corresponding to [log1p](burn_backend::ops::FloatTensorOps::float_log1p).
151    Log1p(UnaryOpIr),
152    /// Operation corresponding to [erf](burn_backend::ops::FloatTensorOps::float_erf).
153    Erf(UnaryOpIr),
154    /// Operation corresponding to [powf_scalar](burn_backend::ops::FloatTensorOps::float_powf_scalar).
155    PowfScalar(ScalarOpIr),
156    /// Operation corresponding to [sqrt](burn_backend::ops::FloatTensorOps::float_sqrt).
157    Sqrt(UnaryOpIr),
158    /// Operation corresponding to [cos](burn_backend::ops::FloatTensorOps::float_cos).
159    Cos(UnaryOpIr),
160    /// Operation corresponding to [cosh](burn_backend::ops::FloatTensorOps::float_cosh).
161    Cosh(UnaryOpIr),
162    /// Operation corresponding to [sin](burn_backend::ops::FloatTensorOps::float_sin).
163    Sin(UnaryOpIr),
164    /// Operation corresponding to [sin](burn_backend::ops::FloatTensorOps::float_sinh).
165    Sinh(UnaryOpIr),
166    /// Operation corresponding to [tan](burn_backend::ops::FloatTensorOps::float_tan).
167    Tan(UnaryOpIr),
168    /// Operation corresponding to [tanh](burn_backend::ops::FloatTensorOps::float_tanh).
169    Tanh(UnaryOpIr),
170    /// Operation corresponding to [acos](burn_backend::ops::FloatTensorOps::float_acos).
171    ArcCos(UnaryOpIr),
172    /// Operation corresponding to [acosh](burn_backend::ops::FloatTensorOps::float_acosh).
173    ArcCosh(UnaryOpIr),
174    /// Operation corresponding to [asin](burn_backend::ops::FloatTensorOps::float_asin).
175    ArcSin(UnaryOpIr),
176    /// Operation corresponding to [asinh](burn_backend::ops::FloatTensorOps::float_asinh).
177    ArcSinh(UnaryOpIr),
178    /// Operation corresponding to [atan](burn_backend::ops::FloatTensorOps::float_atan).
179    ArcTan(UnaryOpIr),
180    /// Operation corresponding to [atanh](burn_backend::ops::FloatTensorOps::float_atanh).
181    ArcTanh(UnaryOpIr),
182    /// Operation corresponding to [atan2](burn_backend::ops::FloatTensorOps::float_atan2).
183    ArcTan2(BinaryOpIr),
184    /// Operation corresponding to [round](burn_backend::ops::FloatTensorOps::float_round).
185    Round(UnaryOpIr),
186    /// Operation corresponding to [floor](burn_backend::ops::FloatTensorOps::float_floor).
187    Floor(UnaryOpIr),
188    /// Operation corresponding to [ceil](burn_backend::ops::FloatTensorOps::float_ceil).
189    Ceil(UnaryOpIr),
190    /// Operation corresponding to [trunc](burn_backend::ops::FloatTensorOps::float_trunc).
191    Trunc(UnaryOpIr),
192    /// Operation corresponding to [into_int](burn_backend::ops::FloatTensorOps::float_into_int).
193    IntoInt(CastOpIr),
194    /// Operation corresponding to [matmul](burn_backend::ops::FloatTensorOps::float_matmul).
195    Matmul(MatmulOpIr),
196    /// Operation corresponding to [cross](burn_backend::ops::FloatTensorOps::float_cross).
197    Cross(CrossOpIr),
198    /// Operation corresponding to [random](burn_backend::ops::FloatTensorOps::float_random).
199    Random(RandomOpIr),
200    /// Operation corresponding to [recip](burn_backend::ops::FloatTensorOps::float_recip).
201    Recip(UnaryOpIr),
202    /// Operation corresponding to [is_nan](burn_backend::ops::FloatTensorOps::float_is_nan).
203    IsNan(UnaryOpIr),
204    /// Operation corresponding to [is_nan](burn_backend::ops::FloatTensorOps::float_is_inf).
205    IsInf(UnaryOpIr),
206    /// Operation corresponding to [quantize](burn_backend::ops::QTensorOps::quantize).
207    Quantize(QuantizeOpIr),
208    /// Operation corresponding to [dequantize](burn_backend::ops::QTensorOps::dequantize).
209    Dequantize(DequantizeOpIr),
210    /// Operation corresponding to [grid_sample_2d](burn_backend::ops::FloatTensorOps::float_grid_sample_2d).
211    GridSample2d(GridSample2dOpIr),
212    /// Operation corresponding to [powf](burn_backend::ops::FloatTensorOps::float_powi).
213    Powf(BinaryOpIr),
214    /// Operation corresponding to [hypot](burn_backend::ops::FloatTensorOps::float_hypot).
215    Hypot(BinaryOpIr),
216    /// Operation corresponding to [fmod](burn_backend::ops::FloatTensorOps::float_fmod).
217    Fmod(BinaryOpIr),
218    /// Operation corresponding to [fmod scalar](burn_backend::ops::FloatTensorOps::float_fmod_scalar).
219    FmodScalar(ScalarOpIr),
220}
221
222/// Operation intermediate representation specific to module.
223#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
224pub enum ModuleOperationIr {
225    /// Batch normalization with explicitly supplied statistics, corresponding
226    /// to [batch_norm](burn_backend::ops::ModuleOps::batch_norm).
227    BatchNorm(BatchNormOpIr),
228    /// Operation corresponding to [embedding](burn_backend::ops::ModuleOps::embedding).
229    Embedding(EmbeddingOpIr),
230    /// Operation corresponding to [embedding_backward](burn_backend::ops::ModuleOps::embedding_backward).
231    EmbeddingBackward(EmbeddingBackwardOpIr),
232    /// Operation corresponding to [linear](burn_backend::ops::ModuleOps::linear).
233    Linear(LinearOpIr),
234    /// Operation corresponding to [linear_x_backward](burn_backend::ops::ModuleOps::linear_x_backward).
235    LinearXBackward(LinearXBackwardOpIr),
236    /// Operation corresponding to [linear_weight_backward](burn_backend::ops::ModuleOps::linear_weight_backward).
237    LinearWeightBackward(LinearWeightBackwardOpIr),
238    /// Operation corresponding to [linear_bias_backward](burn_backend::ops::ModuleOps::linear_bias_backward).
239    LinearBiasBackward(LinearBiasBackwardOpIr),
240    /// Operation corresponding to [conv1d](burn_backend::ops::ModuleOps::conv1d).
241    Conv1d(Conv1dOpIr),
242    /// Operation corresponding to [conv1d_x_backward](burn_backend::ops::ModuleOps::conv1d_x_backward).
243    Conv1dXBackward(Conv1dXBackwardOpIr),
244    /// Operation corresponding to [conv1d_weight_backward](burn_backend::ops::ModuleOps::conv1d_weight_backward).
245    Conv1dWeightBackward(Conv1dWeightBackwardOpIr),
246    /// Operation corresponding to [conv1d_bias_backward](burn_backend::ops::ModuleOps::conv1d_bias_backward).
247    Conv1dBiasBackward(Conv1dBiasBackwardOpIr),
248    /// Operation corresponding to [conv2d](burn_backend::ops::ModuleOps::conv2d).
249    Conv2d(Conv2dOpIr),
250    /// Operation corresponding to [conv2d_x_backward](burn_backend::ops::ModuleOps::conv2d_x_backward).
251    Conv2dXBackward(Conv2dXBackwardOpIr),
252    /// Operation corresponding to [conv2d_weight_backward](burn_backend::ops::ModuleOps::conv2d_weight_backward).
253    Conv2dWeightBackward(Conv2dWeightBackwardOpIr),
254    /// Operation corresponding to [conv2d_bias_backward](burn_backend::ops::ModuleOps::conv2d_bias_backward).
255    Conv2dBiasBackward(Conv2dBiasBackwardOpIr),
256    /// Operation corresponding to [conv3d](burn_backend::ops::ModuleOps::conv3d).
257    Conv3d(Conv3dOpIr),
258    /// Operation corresponding to [conv3d_x_backward](burn_backend::ops::ModuleOps::conv3d_x_backward).
259    Conv3dXBackward(Conv3dXBackwardOpIr),
260    /// Operation corresponding to [conv3d_weight_backward](burn_backend::ops::ModuleOps::conv3d_weight_backward).
261    Conv3dWeightBackward(Conv3dWeightBackwardOpIr),
262    /// Operation corresponding to [conv3d_bias_backward](burn_backend::ops::ModuleOps::conv3d_bias_backward).
263    Conv3dBiasBackward(Conv3dBiasBackwardOpIr),
264    /// Operation corresponding to [deform_conv2d](burn_backend::ops::ModuleOps::deform_conv2d)
265    DeformableConv2d(Box<DeformConv2dOpIr>),
266    /// Operation corresponding to [deform_conv2d_backward](burn_backend::ops::ModuleOps::deform_conv2d_backward)
267    DeformableConv2dBackward(Box<DeformConv2dBackwardOpIr>),
268    /// Operation corresponding to [conv transpose 1d](burn_backend::ops::ModuleOps::conv_transpose1d).
269    ConvTranspose1d(ConvTranspose1dOpIr),
270    /// Operation corresponding to [conv transpose 2d](burn_backend::ops::ModuleOps::conv_transpose2d).
271    ConvTranspose2d(ConvTranspose2dOpIr),
272    /// Operation corresponding to [conv transpose 3d](burn_backend::ops::ModuleOps::conv_transpose3d).
273    ConvTranspose3d(ConvTranspose3dOpIr),
274    /// Operation corresponding to [avg pool 1d](burn_backend::ops::ModuleOps::avg_pool1d).
275    AvgPool1d(AvgPool1dOpIr),
276    /// Operation corresponding to [avg pool 2d](burn_backend::ops::ModuleOps::avg_pool2d).
277    AvgPool2d(AvgPool2dOpIr),
278    /// Operation corresponding to [avg pool 3d](burn_backend::ops::ModuleOps::avg_pool3d).
279    AvgPool3d(AvgPool3dOpIr),
280    /// Operation corresponding to
281    /// [avg pool 1d backward](burn_backend::ops::ModuleOps::avg_pool1d_backward).
282    AvgPool1dBackward(AvgPool1dBackwardOpIr),
283    /// Operation corresponding to
284    /// [avg pool 2d backward](burn_backend::ops::ModuleOps::avg_pool2d_backward).
285    AvgPool2dBackward(AvgPool2dBackwardOpIr),
286    /// Operation corresponding to
287    /// [avg pool 3d backward](burn_backend::ops::ModuleOps::avg_pool3d_backward).
288    AvgPool3dBackward(AvgPool3dBackwardOpIr),
289    /// Operation corresponding to
290    /// [adaptive avg pool 1d](burn_backend::ops::ModuleOps::adaptive_avg_pool1d).
291    AdaptiveAvgPool1d(AdaptiveAvgPool1dOpIr),
292    /// Operation corresponding to
293    /// [adaptive avg pool 2d](burn_backend::ops::ModuleOps::adaptive_avg_pool2d).
294    AdaptiveAvgPool2d(AdaptiveAvgPool2dOpIr),
295    /// Operation corresponding to
296    /// [adaptive avg pool 1d backward](burn_backend::ops::ModuleOps::adaptive_avg_pool1d_backward).
297    AdaptiveAvgPool1dBackward(AdaptiveAvgPool1dBackwardOpIr),
298    /// Operation corresponding to
299    /// [adaptive avg pool 2d backward](burn_backend::ops::ModuleOps::adaptive_avg_pool2d_backward).
300    AdaptiveAvgPool2dBackward(AdaptiveAvgPool2dBackwardOpIr),
301    /// Operation corresponding to
302    /// [adaptive avg pool 3d](burn_backend::ops::ModuleOps::adaptive_avg_pool3d).
303    AdaptiveAvgPool3d(AdaptiveAvgPool3dOpIr),
304    /// Operation corresponding to
305    /// [adaptive avg pool 3d backward](burn_backend::ops::ModuleOps::adaptive_avg_pool3d_backward).
306    AdaptiveAvgPool3dBackward(AdaptiveAvgPool3dBackwardOpIr),
307    /// Operation corresponding to
308    /// [max pool 1d](burn_backend::ops::ModuleOps::max_pool1d).
309    MaxPool1d(MaxPool1dOpIr),
310    /// Operation corresponding to
311    /// [max pool 1d with indices](burn_backend::ops::ModuleOps::max_pool1d_with_indices).
312    MaxPool1dWithIndices(MaxPool1dWithIndicesOpIr),
313    /// Operation corresponding to
314    /// [max pool 1d with indices backward](burn_backend::ops::ModuleOps::max_pool1d_with_indices_backward).
315    MaxPool1dWithIndicesBackward(MaxPool1dWithIndicesBackwardOpIr),
316    /// Operation corresponding to
317    /// [max pool 2d](burn_backend::ops::ModuleOps::max_pool1d).
318    MaxPool2d(MaxPool2dOpIr),
319    /// Operation corresponding to
320    /// [max pool 2d with indices](burn_backend::ops::ModuleOps::max_pool2d_with_indices).
321    MaxPool2dWithIndices(MaxPool2dWithIndicesOpIr),
322    /// Operation corresponding to
323    /// [max pool 2d with indices backward](burn_backend::ops::ModuleOps::max_pool2d_with_indices_backward).
324    MaxPool2dWithIndicesBackward(MaxPool2dWithIndicesBackwardOpIr),
325    /// Operation corresponding to
326    /// [max pool 3d](burn_backend::ops::ModuleOps::max_pool3d).
327    MaxPool3d(MaxPool3dOpIr),
328    /// Operation corresponding to
329    /// [max pool 3d with indices](burn_backend::ops::ModuleOps::max_pool3d_with_indices).
330    MaxPool3dWithIndices(MaxPool3dWithIndicesOpIr),
331    /// Operation corresponding to
332    /// [max pool 3d with indices backward](burn_backend::ops::ModuleOps::max_pool3d_with_indices_backward).
333    MaxPool3dWithIndicesBackward(MaxPool3dWithIndicesBackwardOpIr),
334    /// Operation corresponding to [interpolate](burn_backend::ops::ModuleOps::interpolate).
335    Interpolate(InterpolateOpIr),
336    /// Operation corresponding to [interpolate backward](burn_backend::ops::ModuleOps::interpolate_backward).
337    InterpolateBackward(InterpolateBackwardOpIr),
338    /// Operation corresponding to [attention](burn_backend::ops::ModuleOps::attention).
339    Attention(AttentionOpIr),
340    /// Operation corresponding to
341    /// [attention_with_logsumexp](burn_backend::ops::ModuleOps::attention_with_logsumexp).
342    AttentionWithLogsumexp(AttentionWithLogsumexpOpIr),
343    /// Operation corresponding to
344    /// [attention_backward](burn_backend::ops::ModuleOps::attention_backward).
345    AttentionBackward(Box<AttentionBackwardOpIr>),
346    /// Operation corresponding to [ctc_loss](burn_backend::ops::ModuleOps::ctc_loss).
347    CtcLoss(CtcLossOpIr),
348    /// Operation corresponding to
349    /// [ctc_loss_backward](burn_backend::ops::ModuleOps::ctc_loss_backward).
350    CtcLossBackward(CtcLossBackwardOpIr),
351    /// Operation corresponding to [layer_norm](burn_backend::ops::ModuleOps::layer_norm).
352    LayerNorm(LayerNormOpIr),
353    /// Operation corresponding to [unfold4d](burn_backend::ops::ModuleOps::unfold4d).
354    Unfold4d(Unfold4dOpIr),
355    /// Operation corresponding to
356    /// [conv_transpose1d_weight_backward](burn_backend::ops::ModuleOps::conv_transpose1d_weight_backward).
357    ConvTranspose1dWeightBackward(ConvTranspose1dWeightBackwardOpIr),
358    /// Operation corresponding to
359    /// [conv_transpose1d_bias_backward](burn_backend::ops::ModuleOps::conv_transpose1d_bias_backward).
360    ConvTranspose1dBiasBackward(ConvTranspose1dBiasBackwardOpIr),
361    /// Operation corresponding to
362    /// [conv_transpose2d_weight_backward](burn_backend::ops::ModuleOps::conv_transpose2d_weight_backward).
363    ConvTranspose2dWeightBackward(ConvTranspose2dWeightBackwardOpIr),
364    /// Operation corresponding to
365    /// [conv_transpose2d_bias_backward](burn_backend::ops::ModuleOps::conv_transpose2d_bias_backward).
366    ConvTranspose2dBiasBackward(ConvTranspose2dBiasBackwardOpIr),
367    /// Operation corresponding to
368    /// [conv_transpose3d_weight_backward](burn_backend::ops::ModuleOps::conv_transpose3d_weight_backward).
369    ConvTranspose3dWeightBackward(ConvTranspose3dWeightBackwardOpIr),
370    /// Operation corresponding to
371    /// [conv_transpose3d_bias_backward](burn_backend::ops::ModuleOps::conv_transpose3d_bias_backward).
372    ConvTranspose3dBiasBackward(ConvTranspose3dBiasBackwardOpIr),
373}
374
375/// Basic operations that can be done on any tensor type.
376#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
377pub enum BaseOperationIr {
378    /// Operation corresponding to:
379    ///
380    /// Float => [reshape](burn_backend::ops::FloatTensorOps::float_reshape).
381    /// Int => [reshape](burn_backend::ops::IntTensorOps::int_reshape).
382    /// Bool => [reshape](burn_backend::ops::BoolTensorOps::bool_reshape).
383    Reshape(ShapeOpIr),
384
385    /// Operation corresponding to:
386    ///
387    /// Float => [swap_dims](burn_backend::ops::FloatTensorOps::float_swap_dims).
388    /// Int => [swap_dims](burn_backend::ops::IntTensorOps::int_swap_dims).
389    /// Bool => [swap_dims](burn_backend::ops::BoolTensorOps::bool_swap_dims).
390    SwapDims(SwapDimsOpIr),
391
392    /// Operation corresponding to:
393    ///
394    /// Float => [permute](burn_backend::ops::FloatTensorOps::float_permute).
395    /// Int => [permute](burn_backend::ops::IntTensorOps::int_permute).
396    /// Bool => [permute](burn_backend::ops::BoolTensorOps::bool_permute).
397    Permute(PermuteOpIr),
398
399    /// Operation corresponding to:
400    /// Float => [flip](burn_backend::ops::FloatTensorOps::float_flip).
401    /// Int => [flip](burn_backend::ops::IntTensorOps::int_flip).
402    /// Bool => [flip](burn_backend::ops::BoolTensorOps::bool_flip).
403    Flip(FlipOpIr),
404
405    /// Operation corresponding to:
406    ///
407    /// Float => [expand](burn_backend::ops::FloatTensorOps::float_expand).
408    /// Int => [expand](burn_backend::ops::IntTensorOps::int_expand).
409    /// Bool => [expand](burn_backend::ops::BoolTensorOps::bool_expand).
410    Expand(ShapeOpIr),
411
412    /// Unfold windows along an axis.
413    ///
414    Unfold(UnfoldOpIr),
415
416    /// Operation corresponding to:
417    ///
418    /// Float => [slice](burn_backend::ops::FloatTensorOps::float_slice).
419    /// Int => [slice](burn_backend::ops::IntTensorOps::int_slice).
420    /// Bool => [slice](burn_backend::ops::BoolTensorOps::bool_slice).
421    Slice(SliceOpIr),
422    /// Operation corresponding to:
423    ///
424    /// Float => [slice assign](burn_backend::ops::FloatTensorOps::float_slice_assign).
425    /// Int => [slice assign](burn_backend::ops::IntTensorOps::int_slice_assign).
426    /// Bool => [slice assign](burn_backend::ops::BoolTensorOps::bool_slice_assign).
427    SliceAssign(SliceAssignOpIr),
428    /// Operation corresponding to:
429    ///
430    /// Float => [select](burn_backend::ops::FloatTensorOps::float_select).
431    /// Int => [select](burn_backend::ops::IntTensorOps::int_select).
432    /// Bool => [select](burn_backend::ops::BoolTensorOps::bool_select).
433    Select(SelectOpIr),
434    /// Operation corresponding to:
435    ///
436    /// Float => [select assign](burn_backend::ops::FloatTensorOps::float_select_assign).
437    /// Int => [select assign](burn_backend::ops::IntTensorOps::int_select_assign).
438    /// Bool => [select assign](burn_backend::ops::BoolTensorOps::bool_select_assign).
439    SelectAssign(SelectAssignOpIr),
440    /// Operation corresponding to:
441    ///
442    /// Float => [mask where](burn_backend::ops::FloatTensorOps::float_mask_where).
443    /// Int => [mask where](burn_backend::ops::IntTensorOps::int_mask_where).
444    /// Bool => [mask where](burn_backend::ops::BoolTensorOps::bool_mask_where).
445    MaskWhere(MaskWhereOpIr),
446    /// Operation corresponding to:
447    ///
448    /// Float => [mask fill](burn_backend::ops::FloatTensorOps::float_mask_fill).
449    /// Int => [mask fill](burn_backend::ops::IntTensorOps::int_mask_fill).
450    /// Bool => [mask fill](burn_backend::ops::BoolTensorOps::bool_mask_fill).
451    MaskFill(MaskFillOpIr),
452    /// Operation corresponding to:
453    ///
454    /// Float => [gather](burn_backend::ops::FloatTensorOps::float_gather).
455    /// Int => [gather](burn_backend::ops::IntTensorOps::int_gather).
456    /// Bool => [gather](burn_backend::ops::BoolTensorOps::bool_gather).
457    Gather(GatherOpIr),
458    /// Operation corresponding to:
459    ///
460    /// Float => [scatter](burn_backend::ops::FloatTensorOps::float_scatter).
461    /// Int => [scatter](burn_backend::ops::IntTensorOps::int_scatter).
462    /// Bool => [scatter](burn_backend::ops::BoolTensorOps::bool_scatter).
463    Scatter(ScatterOpIr),
464    /// Multi-dimensional scatter operation.
465    ScatterNd(ScatterNdOpIr),
466    /// Multi-dimensional gather operation.
467    GatherNd(GatherNdOpIr),
468    /// Operation corresponding to:
469    ///
470    /// Float => [equal](burn_backend::ops::FloatTensorOps::float_equal).
471    /// Int => [equal](burn_backend::ops::IntTensorOps::int_equal).
472    /// Bool => [equal](burn_backend::ops::BoolTensorOps::bool_equal).
473    Equal(BinaryOpIr),
474    /// Operation corresponding to:
475    ///
476    /// Float => [equal elem](burn_backend::ops::FloatTensorOps::float_equal_elem).
477    /// Int => [equal elem](burn_backend::ops::IntTensorOps::int_equal_elem).
478    /// Bool => [equal elem](burn_backend::ops::BoolTensorOps::bool_equal_elem).
479    EqualElem(ScalarOpIr),
480    /// Operation corresponding to:
481    ///
482    /// Float => [repeat dim](burn_backend::ops::FloatTensorOps::float_repeat_dim).
483    /// Int => [repeat dim](burn_backend::ops::IntTensorOps::int_repeat_dim).
484    /// Bool => [repeat dim](burn_backend::ops::BoolTensorOps::bool_repeat_dim).
485    RepeatDim(RepeatDimOpIr),
486    /// Operation corresponding to:
487    ///
488    /// Float => [cat](burn_backend::ops::FloatTensorOps::float_cat).
489    /// Int => [cat](burn_backend::ops::IntTensorOps::int_cat).
490    /// Bool => [cat](burn_backend::ops::BoolTensorOps::bool_cat).
491    Cat(CatOpIr),
492    /// Cast operation, no direct operation and should be supported by fusion backend.
493    Cast(CastOpIr),
494    /// Operation corresponding to:
495    ///
496    /// Float => [empty](burn_backend::ops::FloatTensorOps::float_empty).
497    /// Int => [empty](burn_backend::ops::IntTensorOps::int_empty).
498    /// Bool => [empty](burn_backend::ops::BoolTensorOps::bool_empty).
499    Empty(CreationOpIr),
500    /// Operation corresponding to:
501    ///
502    /// Float => [ones](burn_backend::ops::FloatTensorOps::float_ones).
503    /// Int => [ones](burn_backend::ops::IntTensorOps::int_ones).
504    /// Bool => [ones](burn_backend::ops::BoolTensorOps::bool_ones).
505    Ones(CreationOpIr),
506    /// Operation corresponding to:
507    ///
508    /// Float => [zeros](burn_backend::ops::FloatTensorOps::float_zeros).
509    /// Int => [zeros](burn_backend::ops::IntTensorOps::int_zeros).
510    /// Bool => [zeros](burn_backend::ops::BoolTensorOps::bool_zeros).
511    Zeros(CreationOpIr),
512    /// Operation corresponding to:
513    ///
514    /// Float => [not_equal](burn_backend::ops::FloatTensorOps::float_not_equal).
515    /// Int => [not_equal](burn_backend::ops::IntTensorOps::int_not_equal).
516    /// Bool => [not_equal](burn_backend::ops::BoolTensorOps::bool_not_equal).
517    NotEqual(BinaryOpIr),
518    /// Operation corresponding to:
519    ///
520    /// Float => [not_equal_elem](burn_backend::ops::FloatTensorOps::float_not_equal_elem).
521    /// Int => [not_equal_elem](burn_backend::ops::IntTensorOps::int_not_equal_elem).
522    /// Bool => [not_equal_elem](burn_backend::ops::BoolTensorOps::bool_not_equal_elem).
523    NotEqualElem(ScalarOpIr),
524    /// Reduce-`all` over the input tensor.
525    ///
526    /// Float/Int input is treated as a non-zero check; bool input is the value itself.
527    /// Output is a single-element bool tensor.
528    All(ReduceOpIr),
529    /// Reduce-`any` over the input tensor.
530    Any(ReduceOpIr),
531    /// Reduce-`all` along a dim.
532    AllDim(ReduceDimOpIr),
533    /// Reduce-`any` along a dim.
534    AnyDim(ReduceDimOpIr),
535}
536
537/// Numeric operations on int and float tensors.
538#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
539pub enum NumericOperationIr {
540    /// Operation corresponding to:
541    ///
542    /// Float => [add](burn_backend::ops::FloatTensorOps::float_add).
543    /// Int => [add](burn_backend::ops::IntTensorOps::int_add).
544    Add(BinaryOpIr),
545    /// Operation corresponding to:
546    ///
547    /// Float => [add scalar](burn_backend::ops::FloatTensorOps::float_add_scalar).
548    /// Int => [add scalar](burn_backend::ops::IntTensorOps::int_add_scalar).
549    AddScalar(ScalarOpIr),
550    /// Operation corresponding to:
551    ///
552    /// Float => [sub](burn_backend::ops::FloatTensorOps::float_sub).
553    /// Int => [sub](burn_backend::ops::IntTensorOps::int_sub).
554    Sub(BinaryOpIr),
555    /// Operation corresponding to:
556    ///
557    /// Float => [sub scalar](burn_backend::ops::FloatTensorOps::float_sub_scalar).
558    /// Int => [sub scalar](burn_backend::ops::IntTensorOps::int_sub_scalar).
559    SubScalar(ScalarOpIr),
560    /// Operation corresponding to:
561    ///
562    /// Float => [div](burn_backend::ops::FloatTensorOps::float_div).
563    /// Int => [div](burn_backend::ops::IntTensorOps::int_div).
564    Div(BinaryOpIr),
565    /// Operation corresponding to:
566    ///
567    /// Float => [div scalar](burn_backend::ops::FloatTensorOps::float_div_scalar).
568    /// Int => [div scalar](burn_backend::ops::IntTensorOps::int_div_scalar).
569    DivScalar(ScalarOpIr),
570    /// Operation corresponding to:
571    ///
572    /// Float => [rem](burn_backend::ops::FloatTensorOps::float_remainder).
573    /// Int => [rem](burn_backend::ops::IntTensorOps::int_remainder).
574    Rem(BinaryOpIr),
575    /// Operation corresponding to:
576    ///
577    /// Float => [rem scalar](burn_backend::ops::FloatTensorOps::float_remainder_scalar).
578    /// Int => [rem scalar](burn_backend::ops::IntTensorOps::int_remainder_scalar).
579    RemScalar(ScalarOpIr),
580    /// Operation corresponding to:
581    ///
582    /// Float => [mul](burn_backend::ops::FloatTensorOps::float_mul).
583    /// Int => [mul](burn_backend::ops::IntTensorOps::int_mul).
584    Mul(BinaryOpIr),
585    /// Operation corresponding to:
586    ///
587    /// Float => [mul scalar](burn_backend::ops::FloatTensorOps::float_mul_scalar).
588    /// Int => [mul scalar](burn_backend::ops::IntTensorOps::int_mul_scalar).
589    MulScalar(ScalarOpIr),
590    /// Operation corresponding to:
591    ///
592    /// Float => [abs](burn_backend::ops::FloatTensorOps::float_abs).
593    /// Int => [abs](burn_backend::ops::IntTensorOps::int_abs).
594    Abs(UnaryOpIr),
595    /// Operation corresponding to:
596    ///
597    /// Float => [full](burn_backend::ops::FloatTensorOps::float_full).
598    /// Int => [full](burn_backend::ops::IntTensorOps::int_full).
599    Full(FullOpIr),
600    /// Operation corresponding to:
601    ///
602    /// Float => [mean dim](burn_backend::ops::FloatTensorOps::float_mean_dim).
603    /// Int => [mean dim](burn_backend::ops::IntTensorOps::int_mean_dim).
604    MeanDim(ReduceDimOpIr),
605    /// Operation corresponding to:
606    ///
607    /// Float => [mean](burn_backend::ops::FloatTensorOps::float_mean).
608    /// Int => [mean](burn_backend::ops::IntTensorOps::int_mean).
609    Mean(ReduceOpIr),
610    /// Operation corresponding to:
611    ///
612    /// Float => [sum](burn_backend::ops::FloatTensorOps::float_sum).
613    /// Int => [sum](burn_backend::ops::IntTensorOps::int_sum).
614    Sum(ReduceOpIr),
615    /// Operation corresponding to:
616    ///
617    /// Float => [sum dim](burn_backend::ops::FloatTensorOps::float_sum_dim).
618    /// Int => [sum dim](burn_backend::ops::IntTensorOps::int_sum_dim).
619    SumDim(ReduceDimOpIr),
620    /// Operation corresponding to summing several dimensions at once.
621    SumDims(ReduceDimsOpIr),
622    /// Operation corresponding to:
623    ///
624    /// Float => [prod](burn_backend::ops::FloatTensorOps::float_prod).
625    /// Int => [prod](burn_backend::ops::IntTensorOps::int_prod).
626    Prod(ReduceOpIr),
627    /// Operation corresponding to:
628    ///
629    /// Float => [prod dim](burn_backend::ops::FloatTensorOps::float_prod_dim).
630    /// Int => [prod dim](burn_backend::ops::IntTensorOps::int_prod_dim).
631    ProdDim(ReduceDimOpIr),
632    /// Operation corresponding to:
633    ///
634    /// Float => [greater](burn_backend::ops::FloatTensorOps::float_greater).
635    /// Int => [greater](burn_backend::ops::IntTensorOps::int_greater).
636    Greater(BinaryOpIr),
637    /// Operation corresponding to:
638    ///
639    /// Float => [greater elem](burn_backend::ops::FloatTensorOps::float_greater_elem).
640    /// Int => [greater elem](burn_backend::ops::IntTensorOps::int_greater_elem).
641    GreaterElem(ScalarOpIr),
642    /// Operation corresponding to:
643    ///
644    /// Float => [greater equal](burn_backend::ops::FloatTensorOps::float_greater_elem).
645    /// Int => [greater elem](burn_backend::ops::IntTensorOps::int_greater_elem).
646    GreaterEqual(BinaryOpIr),
647    /// Operation corresponding to:
648    ///
649    /// Float => [greater equal elem](burn_backend::ops::FloatTensorOps::float_greater_equal_elem).
650    /// Int => [greater equal elem](burn_backend::ops::IntTensorOps::int_greater_equal_elem).
651    GreaterEqualElem(ScalarOpIr),
652    /// Operation corresponding to:
653    ///
654    /// Float => [lower](burn_backend::ops::FloatTensorOps::float_lower).
655    /// Int => [lower](burn_backend::ops::IntTensorOps::int_lower).
656    Lower(BinaryOpIr),
657    /// Operation corresponding to:
658    ///
659    /// Float => [lower elem](burn_backend::ops::FloatTensorOps::float_lower_elem).
660    /// Int => [lower elem](burn_backend::ops::IntTensorOps::int_lower_elem).
661    LowerElem(ScalarOpIr),
662    /// Operation corresponding to:
663    ///
664    /// Float => [lower equal](burn_backend::ops::FloatTensorOps::float_lower_equal).
665    /// Int => [lower equal](burn_backend::ops::IntTensorOps::int_lower_equal).
666    LowerEqual(BinaryOpIr),
667    /// Operation corresponding to:
668    ///
669    /// Float => [lower equal elem](burn_backend::ops::FloatTensorOps::float_lower_equal_elem).
670    /// Int => [lower equal elem](burn_backend::ops::IntTensorOps::int_lower_equal_elem).
671    LowerEqualElem(ScalarOpIr),
672    /// Operation corresponding to:
673    ///
674    /// Float => [argmax](burn_backend::ops::FloatTensorOps::float_argmax).
675    /// Int => [argmax](burn_backend::ops::IntTensorOps::int_argmax).
676    ArgMax(ReduceDimOpIr),
677    /// Operation corresponding to:
678    ///
679    /// Float => [argtopk](burn_backend::ops::FloatTensorOps::float_argtopk).
680    /// Int => [argtopk](burn_backend::ops::IntTensorOps::int_argtopk).
681    ArgTopK(ReduceDimOpIr),
682    /// Operation corresponding to:
683    ///
684    /// Float => [topk](burn_backend::ops::FloatTensorOps::float_topk).
685    /// Int => [topk](burn_backend::ops::IntTensorOps::int_topk).
686    TopK(ReduceDimOpIr),
687    /// Operation corresponding to:
688    ///
689    /// Float => [topk with indices](burn_backend::ops::FloatTensorOps::float_topk_with_indices).
690    /// Int => [topk with indices](burn_backend::ops::IntTensorOps::int_topk_with_indices).
691    TopKWithIndices(TopKWithIndicesOpIr),
692    /// Operation corresponding to:
693    ///
694    /// Float => [argmin](burn_backend::ops::FloatTensorOps::float_argmin).
695    /// Int => [argmin](burn_backend::ops::IntTensorOps::int_argmin).
696    ArgMin(ReduceDimOpIr),
697    /// Operation corresponding to:
698    ///
699    /// Float => [max](burn_backend::ops::FloatTensorOps::float_max).
700    /// Int => [max](burn_backend::ops::IntTensorOps::int_max).
701    Max(ReduceOpIr),
702    /// Operation corresponding to:
703    ///
704    /// Float => [max dim with indices](burn_backend::ops::FloatTensorOps::float_max_dim_with_indices).
705    /// Int => [max dim with indices](burn_backend::ops::IntTensorOps::int_max_dim_with_indices).
706    MaxDimWithIndices(ReduceDimWithIndicesOpIr),
707    /// Operation corresponding to:
708    ///
709    /// Float => [min dim with indices](burn_backend::ops::FloatTensorOps::float_min_dim_with_indices).
710    /// Int => [min dim with indices](burn_backend::ops::IntTensorOps::int_min_dim_with_indices).
711    MinDimWithIndices(ReduceDimWithIndicesOpIr),
712    /// Operation corresponding to:
713    ///
714    /// Float => [min](burn_backend::ops::FloatTensorOps::float_min).
715    /// Int => [min](burn_backend::ops::IntTensorOps::int_min).
716    Min(ReduceOpIr),
717    /// Operation corresponding to:
718    ///
719    /// Float => [max dim](burn_backend::ops::FloatTensorOps::float_max_dim).
720    /// Int => [max dim](burn_backend::ops::IntTensorOps::int_max_dim).
721    MaxDim(ReduceDimOpIr),
722    /// Operation corresponding to:
723    ///
724    /// Float => [min dim](burn_backend::ops::FloatTensorOps::float_min_dim).
725    /// Int => [min dim](burn_backend::ops::IntTensorOps::int_min_dim).
726    MinDim(ReduceDimOpIr),
727    /// Operation corresponding to:
728    ///
729    /// Float => [max_abs](burn_backend::ops::FloatTensorOps::float_max_abs).
730    /// Int => [max_abs](burn_backend::ops::IntTensorOps::int_max_abs).
731    MaxAbs(ReduceOpIr),
732    /// Operation corresponding to:
733    ///
734    /// Float => [max_abs dim](burn_backend::ops::FloatTensorOps::float_max_abs_dim).
735    /// Int => [max_abs dim](burn_backend::ops::IntTensorOps::int_max_abs_dim).
736    MaxAbsDim(ReduceDimOpIr),
737    /// Operation corresponding to:
738    ///
739    /// Float => [clamp](burn_backend::ops::FloatTensorOps::float_clamp).
740    /// Int => [clamp](burn_backend::ops::IntTensorOps::int_clamp).
741    Clamp(ClampOpIr),
742    /// Operation corresponding to:
743    ///
744    /// Int => [random](burn_backend::ops::IntTensorOps::int_random).
745    IntRandom(RandomOpIr),
746    /// Operation corresponding to:
747    ///
748    /// Float => [powf](burn_backend::ops::FloatTensorOps::float_powi).
749    /// Int => [powf](burn_backend::ops::IntTensorOps::int_powi).
750    Powi(BinaryOpIr),
751    /// Operation corresponding to:
752    ///
753    /// Float => [powi_scalar](burn_backend::ops::FloatTensorOps::float_powi_scalar).
754    /// Int => [powi_scalar](burn_backend::ops::IntTensorOps::int_powi_scalar).
755    PowiScalar(ScalarOpIr),
756    /// Operation corresponding to:
757    ///
758    /// Float => [cumsum](burn_backend::ops::FloatTensorOps::float_cumsum).
759    /// Int => [cumsum](burn_backend::ops::IntTensorOps::int_cumsum).
760    CumSum(DimOpIr),
761    /// Operation corresponding to:
762    ///
763    /// Float => [cumprod](burn_backend::ops::FloatTensorOps::float_cumprod).
764    /// Int => [cumprod](burn_backend::ops::IntTensorOps::int_cumprod).
765    CumProd(DimOpIr),
766    /// Operation corresponding to:
767    ///
768    /// Float => [cummin](burn_backend::ops::FloatTensorOps::float_cummin).
769    /// Int => [cummin](burn_backend::ops::IntTensorOps::int_cummin).
770    CumMin(DimOpIr),
771    /// Operation corresponding to:
772    ///
773    /// Float => [cummax](burn_backend::ops::FloatTensorOps::float_cummax).
774    /// Int => [cummax](burn_backend::ops::IntTensorOps::int_cummax).
775    CumMax(DimOpIr),
776    /// Operation corresponding to:
777    ///
778    /// Float => [neg](burn_backend::ops::FloatTensorOps::float_neg).
779    /// Int => [neg](burn_backend::ops::IntTensorOps::int_neg).
780    Neg(UnaryOpIr),
781    /// Operation corresponding to:
782    ///
783    /// Float => [sign](burn_backend::ops::FloatTensorOps::float_sign).
784    /// Int => [sign](burn_backend::ops::IntTensorOps::int_sign).
785    Sign(UnaryOpIr),
786    /// Operation corresponding to:
787    ///
788    /// Float => [clamp_min](burn_backend::ops::FloatTensorOps::float_clamp_min).
789    /// Int => [clamp_min](burn_backend::ops::IntTensorOps::int_clamp_min).
790    ClampMin(ScalarOpIr),
791    /// Operation corresponding to:
792    ///
793    /// Float => [clamp_max](burn_backend::ops::FloatTensorOps::float_clamp_max).
794    /// Int => [clamp_max](burn_backend::ops::IntTensorOps::int_clamp_max).
795    ClampMax(ScalarOpIr),
796    /// Sort along a dim. Output shares the input's shape and dtype.
797    Sort(SortOpIr),
798    /// Sort along a dim, also returning the source indices.
799    SortWithIndices(SortWithIndicesOpIr),
800    /// Sort along a dim and return only the indices (i.e., the argsort).
801    ///
802    /// Shares [`SortOpIr`] with [`Sort`](Self::Sort); only the output dtype differs.
803    ArgSort(SortOpIr),
804    /// Operation corresponding to:
805    ///
806    /// Float => [pad](burn_backend::ops::FloatTensorOps::float_pad).
807    /// Int => [pad](burn_backend::ops::IntTensorOps::int_pad).
808    Pad(PadOpIr),
809}
810
811/// Operation intermediate representation specific to an int tensor.
812#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
813pub enum IntOperationIr {
814    /// Operation corresponding to [into float](burn_backend::ops::IntTensorOps::int_into_float).
815    IntoFloat(CastOpIr),
816    /// Operation corresponding to:
817    ///
818    /// Int => [bitwise and](burn_backend::ops::IntTensorOps::bitwise_and).
819    BitwiseAnd(BinaryOpIr),
820    /// Operation corresponding to:
821    ///
822    /// Int => [bitwise and scalar](burn_backend::ops::IntTensorOps::bitwise_and_scalar).
823    BitwiseAndScalar(ScalarOpIr),
824    /// Operation corresponding to:
825    ///
826    /// Int => [bitwise or](burn_backend::ops::IntTensorOps::bitwise_or).
827    BitwiseOr(BinaryOpIr),
828    /// Operation corresponding to:
829    ///
830    /// Int => [bitwise or scalar](burn_backend::ops::IntTensorOps::bitwise_or_scalar).
831    BitwiseOrScalar(ScalarOpIr),
832    /// Operation corresponding to:
833    ///
834    /// Int => [bitwise xor](burn_backend::ops::IntTensorOps::bitwise_xor).
835    BitwiseXor(BinaryOpIr),
836    /// Operation corresponding to:
837    ///
838    /// Int => [bitwise xor scalar](burn_backend::ops::IntTensorOps::bitwise_xor_scalar).
839    BitwiseXorScalar(ScalarOpIr),
840    /// Operation corresponding to:
841    ///
842    /// Int => [bitwise not](burn_backend::ops::IntTensorOps::bitwise_not).
843    BitwiseNot(UnaryOpIr),
844    /// Operation corresponding to:
845    ///
846    /// Int => [bitwise left shift](burn_backend::ops::IntTensorOps::bitwise_left_shift).
847    BitwiseLeftShift(BinaryOpIr),
848    /// Operation corresponding to:
849    ///
850    /// Int => [bitwise left shift scalar](burn_backend::ops::IntTensorOps::bitwise_left_shift_scalar).
851    BitwiseLeftShiftScalar(ScalarOpIr),
852    /// Operation corresponding to:
853    ///
854    /// Int => [bitwise right shift](burn_backend::ops::IntTensorOps::bitwise_right_shift).
855    BitwiseRightShift(BinaryOpIr),
856    /// Operation corresponding to:
857    ///
858    /// Int => [bitwise right shift scalar](burn_backend::ops::IntTensorOps::bitwise_right_shift_scalar).
859    BitwiseRightShiftScalar(ScalarOpIr),
860    /// Operation corresponding to [matmul](burn_backend::ops::IntTensorOps::int_matmul).
861    Matmul(MatmulOpIr),
862}
863
864/// Operation intermediate representation specific to a bool tensor.
865#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
866pub enum BoolOperationIr {
867    /// Operation corresponding to [into float](burn_backend::ops::BoolTensorOps::bool_into_float).
868    IntoFloat(CastOpIr),
869    /// Operation corresponding to [into int](burn_backend::ops::BoolTensorOps::bool_into_int).
870    IntoInt(CastOpIr),
871    /// Operation corresponding to [not](burn_backend::ops::BoolTensorOps::bool_not).
872    Not(UnaryOpIr),
873    /// Operation corresponding to [and](burn_backend::ops::BoolTensorOps::bool_and).
874    And(BinaryOpIr),
875    /// Operation corresponding to [or](burn_backend::ops::BoolTensorOps::bool_or).
876    Or(BinaryOpIr),
877    /// Operation corresponding to [xor](burn_backend::ops::BoolTensorOps::bool_xor).
878    Xor(BinaryOpIr),
879}
880
881/// Swap dim operation intermediate representation.
882#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
883pub struct SwapDimsOpIr {
884    /// Input tensor intermediate representation.
885    pub input: TensorIr,
886    /// Output tensor intermediate representation.
887    pub out: TensorIr,
888    /// The first dim to swap.
889    pub dim1: usize,
890    /// The second dim to swap.
891    pub dim2: usize,
892}
893
894/// Permute operation intermediate representation.
895#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
896pub struct PermuteOpIr {
897    /// Input tensor intermediate representation.
898    pub input: TensorIr,
899    /// Output tensor intermediate representation.
900    pub out: TensorIr,
901    /// The new order of the dimensions.
902    pub axes: Vec<usize>,
903}
904
905/// Shape operation intermediate representation.
906#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
907pub struct ShapeOpIr {
908    /// Input tensor intermediate representation.
909    pub input: TensorIr,
910    /// Output tensor intermediate representation with the new shape.
911    pub out: TensorIr,
912}
913
914/// Unfold operation intermediate representation.
915#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
916pub struct UnfoldOpIr {
917    /// Input tensor intermediate representation.
918    pub input: TensorIr,
919    /// Output tensor intermediate representation.
920    pub out: TensorIr,
921
922    /// The selected dim.
923    pub dim: usize,
924    /// The window size.
925    pub size: usize,
926    /// The window step along dim.
927    pub step: usize,
928}
929
930/// Flip operation intermediate representation.
931#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
932pub struct FlipOpIr {
933    /// Input tensor intermediate representation.
934    pub input: TensorIr,
935    /// Output tensor intermediate representation.
936    pub out: TensorIr,
937    /// The dimensions to flip.
938    pub axes: Vec<usize>,
939}
940
941/// Padding operation intermediate representation.
942#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
943pub struct PadOpIr {
944    /// Input tensor intermediate representation.
945    pub input: TensorIr,
946    /// Output tensor intermediate representation.
947    pub out: TensorIr,
948    /// One `(before, after)` padding pair per dimension.
949    pub padding: Vec<(usize, usize)>,
950    /// Padding mode.
951    pub mode: PadModeIr,
952}
953
954/// Serializable padding mode intermediate representation.
955#[derive(Clone, Copy, Debug, Hash, PartialEq, Serialize, Deserialize)]
956pub enum PadModeIr {
957    /// Fill padded regions with a constant value.
958    Constant(ScalarIr),
959    /// Reflect values at the boundary, excluding the edge value.
960    Reflect,
961    /// Replicate boundary values.
962    Edge,
963}
964
965impl From<PadMode> for PadModeIr {
966    fn from(value: PadMode) -> Self {
967        match value {
968            PadMode::Constant(value) => Self::Constant(ScalarIr::Float(value as f64)),
969            PadMode::Reflect => Self::Reflect,
970            PadMode::Edge => Self::Edge,
971        }
972    }
973}
974
975impl From<PadModeIr> for PadMode {
976    fn from(value: PadModeIr) -> Self {
977        match value {
978            PadModeIr::Constant(value) => Self::Constant(value.elem()),
979            PadModeIr::Reflect => Self::Reflect,
980            PadModeIr::Edge => Self::Edge,
981        }
982    }
983}
984
985#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
986#[allow(missing_docs)]
987pub struct RandomOpIr {
988    pub out: TensorIr,
989    pub distribution: Distribution,
990}
991
992/// Creation operation intermediate representation.
993/// As opposed to [InitOperationIr], creation operations are lazy initialized.
994#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
995pub struct CreationOpIr {
996    /// Output tensor intermediate representation.
997    pub out: TensorIr,
998}
999
1000/// Full operation intermediate representation.
1001#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1002pub struct FullOpIr {
1003    /// Output tensor intermediate representation.
1004    pub out: TensorIr,
1005    /// Fill value.
1006    pub value: ScalarIr,
1007}
1008
1009#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1010/// Declares a tensor has been initialized.
1011///
1012/// It is necessary to register for proper orphan detection and avoid memory leak.
1013pub struct InitOperationIr {
1014    /// The initialized tensor.
1015    pub out: TensorIr,
1016}
1017
1018#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1019#[allow(missing_docs)]
1020pub struct BinaryOpIr {
1021    pub lhs: TensorIr,
1022    pub rhs: TensorIr,
1023    pub out: TensorIr,
1024}
1025
1026#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1027#[allow(missing_docs)]
1028pub struct MatmulOpIr {
1029    pub lhs: TensorIr,
1030    pub rhs: TensorIr,
1031    pub out: TensorIr,
1032}
1033
1034#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1035#[allow(missing_docs)]
1036pub struct CrossOpIr {
1037    pub lhs: TensorIr,
1038    pub rhs: TensorIr,
1039    pub out: TensorIr,
1040    pub dim: usize,
1041}
1042
1043#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1044#[allow(missing_docs)]
1045pub struct UnaryOpIr {
1046    pub input: TensorIr,
1047    pub out: TensorIr,
1048}
1049
1050#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1051#[allow(missing_docs)]
1052pub struct ScalarOpIr {
1053    pub lhs: TensorIr,
1054    // TODO: Make that an enum with `Value` and `Id` variants for relative/global
1055    // conversion.
1056    pub rhs: ScalarIr,
1057    pub out: TensorIr,
1058}
1059
1060#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Hash)]
1061#[allow(missing_docs)]
1062pub struct ReduceOpIr {
1063    pub input: TensorIr,
1064    pub out: TensorIr,
1065}
1066
1067#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Hash)]
1068#[allow(missing_docs)]
1069pub struct ReduceDimOpIr {
1070    pub input: TensorIr,
1071    pub out: TensorIr,
1072    pub axis: usize,
1073    pub accumulator_len: usize,
1074}
1075
1076/// A reduction over several dimensions at once, each kept with length one.
1077#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Hash)]
1078#[allow(missing_docs)]
1079pub struct ReduceDimsOpIr {
1080    pub input: TensorIr,
1081    pub out: TensorIr,
1082    pub axes: Vec<usize>,
1083}
1084
1085#[allow(missing_docs)]
1086impl ReduceDimsOpIr {
1087    pub fn create<F>(input: TensorIr, axes: Vec<usize>, mut new_id: F) -> Self
1088    where
1089        F: FnMut() -> crate::TensorId,
1090    {
1091        let mut shape = input.shape.clone();
1092        for axis in &axes {
1093            shape[*axis] = 1;
1094        }
1095        let dtype = input.dtype;
1096
1097        Self {
1098            out: TensorIr::uninit(new_id(), shape, dtype),
1099            input,
1100            axes,
1101        }
1102    }
1103}
1104
1105#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1106#[allow(missing_docs)]
1107pub struct CastOpIr {
1108    pub input: TensorIr,
1109    pub out: TensorIr,
1110}
1111
1112/// IR for operations that operate along a dimension without reducing it.
1113/// Unlike `ReduceDimOpIr`, the output shape is the same as the input shape.
1114#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Hash)]
1115#[allow(missing_docs)]
1116pub struct DimOpIr {
1117    pub input: TensorIr,
1118    pub out: TensorIr,
1119    pub axis: usize,
1120}
1121
1122#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1123#[allow(missing_docs)]
1124pub struct GatherOpIr {
1125    pub tensor: TensorIr,
1126    pub dim: usize,
1127    pub indices: TensorIr,
1128    pub out: TensorIr,
1129}
1130
1131#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1132#[allow(missing_docs)]
1133pub struct ScatterOpIr {
1134    pub tensor: TensorIr,
1135    pub dim: usize,
1136    pub indices: TensorIr,
1137    pub value: TensorIr,
1138    pub update: IndexingUpdateOp,
1139    pub out: TensorIr,
1140}
1141
1142#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1143#[allow(missing_docs)]
1144pub struct ScatterNdOpIr {
1145    pub data: TensorIr,
1146    pub indices: TensorIr,
1147    pub values: TensorIr,
1148    pub reduction: IndexingUpdateOp,
1149    pub out: TensorIr,
1150}
1151
1152#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1153#[allow(missing_docs)]
1154pub struct GatherNdOpIr {
1155    pub data: TensorIr,
1156    pub indices: TensorIr,
1157    pub out: TensorIr,
1158}
1159
1160#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1161#[allow(missing_docs)]
1162pub struct SelectOpIr {
1163    pub tensor: TensorIr,
1164    pub dim: usize,
1165    pub indices: TensorIr,
1166    pub out: TensorIr,
1167}
1168
1169#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1170#[allow(missing_docs)]
1171pub struct SelectAssignOpIr {
1172    pub tensor: TensorIr,
1173    pub dim: usize,
1174    pub indices: TensorIr,
1175    pub value: TensorIr,
1176    pub update: IndexingUpdateOp,
1177    pub out: TensorIr,
1178}
1179
1180#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1181#[allow(missing_docs)]
1182pub struct SliceOpIr {
1183    pub tensor: TensorIr,
1184    pub ranges: Vec<Slice>,
1185    pub out: TensorIr,
1186}
1187
1188#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1189#[allow(missing_docs)]
1190pub struct SliceAssignOpIr {
1191    pub tensor: TensorIr,
1192    pub ranges: Vec<burn_backend::Slice>,
1193    pub value: TensorIr,
1194    pub out: TensorIr,
1195}
1196
1197#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1198#[allow(missing_docs)]
1199pub struct MaskWhereOpIr {
1200    pub tensor: TensorIr,
1201    pub mask: TensorIr,
1202    pub value: TensorIr,
1203    pub out: TensorIr,
1204}
1205
1206#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1207#[allow(missing_docs)]
1208pub struct MaskFillOpIr {
1209    pub tensor: TensorIr,
1210    pub mask: TensorIr,
1211    pub value: ScalarIr,
1212    pub out: TensorIr,
1213}
1214
1215#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1216#[allow(missing_docs)]
1217pub struct ClampOpIr {
1218    pub tensor: TensorIr,
1219    pub min: ScalarIr,
1220    pub max: ScalarIr,
1221    pub out: TensorIr,
1222}
1223
1224#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1225#[allow(missing_docs)]
1226pub struct RepeatDimOpIr {
1227    pub tensor: TensorIr,
1228    pub dim: usize,
1229    pub times: usize,
1230    pub out: TensorIr,
1231}
1232
1233#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1234#[allow(missing_docs)]
1235pub struct CatOpIr {
1236    pub tensors: Vec<TensorIr>,
1237    pub dim: usize,
1238    pub out: TensorIr,
1239}
1240
1241#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1242#[allow(missing_docs)]
1243pub struct ReduceDimWithIndicesOpIr {
1244    pub tensor: TensorIr,
1245    pub dim: usize,
1246    pub out: TensorIr,
1247    pub out_indices: TensorIr,
1248}
1249
1250#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1251#[allow(missing_docs)]
1252/// Like [`ReduceDimWithIndicesOpIr`], but for a top-k: the reduced axis keeps `k` entries
1253/// instead of collapsing to 1, so `k` has to be carried explicitly.
1254pub struct TopKWithIndicesOpIr {
1255    pub tensor: TensorIr,
1256    pub dim: usize,
1257    pub k: usize,
1258    pub out: TensorIr,
1259    pub out_indices: TensorIr,
1260}
1261
1262#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1263#[allow(missing_docs)]
1264pub struct EmbeddingOpIr {
1265    pub weights: TensorIr,
1266    pub indices: TensorIr,
1267    pub out: TensorIr,
1268}
1269
1270#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1271#[allow(missing_docs)]
1272pub struct EmbeddingBackwardOpIr {
1273    pub weights: TensorIr,
1274    pub out_grad: TensorIr,
1275    pub indices: TensorIr,
1276    pub out: TensorIr,
1277}
1278
1279/// Batch normalization using explicitly supplied channel statistics.
1280///
1281/// This operation neither calculates nor updates the supplied statistics.
1282#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1283#[allow(missing_docs)]
1284pub struct BatchNormOpIr {
1285    pub x: TensorIr,
1286    pub gamma: TensorIr,
1287    pub beta: TensorIr,
1288    pub mean: TensorIr,
1289    pub variance: TensorIr,
1290    pub epsilon: ScalarIr,
1291    pub out: TensorIr,
1292}
1293
1294#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1295#[allow(missing_docs)]
1296pub struct LinearOpIr {
1297    pub x: TensorIr,
1298    pub weight: TensorIr,
1299    pub bias: Option<TensorIr>,
1300    pub out: TensorIr,
1301}
1302
1303#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1304#[allow(missing_docs)]
1305pub struct LinearXBackwardOpIr {
1306    pub weight: TensorIr,
1307    pub output_grad: TensorIr,
1308    pub out: TensorIr,
1309}
1310
1311#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1312#[allow(missing_docs)]
1313pub struct LinearWeightBackwardOpIr {
1314    pub x: TensorIr,
1315    pub output_grad: TensorIr,
1316    pub out: TensorIr,
1317}
1318
1319#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1320#[allow(missing_docs)]
1321pub struct LinearBiasBackwardOpIr {
1322    pub output_grad: TensorIr,
1323    pub out: TensorIr,
1324}
1325
1326#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1327#[allow(missing_docs)]
1328pub struct Conv1dOpIr {
1329    pub x: TensorIr,
1330    pub weight: TensorIr,
1331    pub bias: Option<TensorIr>,
1332    pub options: Conv1dOptionsIr,
1333    pub out: TensorIr,
1334}
1335
1336#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1337#[allow(missing_docs)]
1338pub struct Conv1dXBackwardOpIr {
1339    pub x: TensorIr,
1340    pub weight: TensorIr,
1341    pub output_grad: TensorIr,
1342    pub options: Conv1dOptionsIr,
1343    pub out: TensorIr,
1344}
1345
1346#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1347#[allow(missing_docs)]
1348pub struct Conv1dWeightBackwardOpIr {
1349    pub x: TensorIr,
1350    pub weight: TensorIr,
1351    pub output_grad: TensorIr,
1352    pub options: Conv1dOptionsIr,
1353    pub out: TensorIr,
1354}
1355
1356#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1357#[allow(missing_docs)]
1358pub struct Conv1dBiasBackwardOpIr {
1359    pub x: TensorIr,
1360    pub bias: TensorIr,
1361    pub output_grad: TensorIr,
1362    pub out: TensorIr,
1363}
1364
1365#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1366#[allow(missing_docs)]
1367pub struct Conv2dOpIr {
1368    pub x: TensorIr,
1369    pub weight: TensorIr,
1370    pub bias: Option<TensorIr>,
1371    pub options: Conv2dOptionsIr,
1372    pub out: TensorIr,
1373}
1374
1375#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1376#[allow(missing_docs)]
1377pub struct Conv2dXBackwardOpIr {
1378    pub x: TensorIr,
1379    pub weight: TensorIr,
1380    pub output_grad: TensorIr,
1381    pub options: Conv2dOptionsIr,
1382    pub out: TensorIr,
1383}
1384
1385#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1386#[allow(missing_docs)]
1387pub struct Conv2dWeightBackwardOpIr {
1388    pub x: TensorIr,
1389    pub weight: TensorIr,
1390    pub output_grad: TensorIr,
1391    pub options: Conv2dOptionsIr,
1392    pub out: TensorIr,
1393}
1394
1395#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1396#[allow(missing_docs)]
1397pub struct Conv2dBiasBackwardOpIr {
1398    pub x: TensorIr,
1399    pub bias: TensorIr,
1400    pub output_grad: TensorIr,
1401    pub out: TensorIr,
1402}
1403
1404#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1405#[allow(missing_docs)]
1406pub struct DeformConv2dOpIr {
1407    pub x: TensorIr,
1408    pub offset: TensorIr,
1409    pub weight: TensorIr,
1410    pub mask: Option<TensorIr>,
1411    pub bias: Option<TensorIr>,
1412    pub options: DeformableConv2dOptionsIr,
1413    pub out: TensorIr,
1414}
1415
1416#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1417#[allow(missing_docs)]
1418pub struct DeformConv2dBackwardOpIr {
1419    pub x: TensorIr,
1420    pub offset: TensorIr,
1421    pub weight: TensorIr,
1422    pub mask: Option<TensorIr>,
1423    pub bias: Option<TensorIr>,
1424    pub out_grad: TensorIr,
1425    pub options: DeformableConv2dOptionsIr,
1426    pub input_grad: TensorIr,
1427    pub offset_grad: TensorIr,
1428    pub weight_grad: TensorIr,
1429    pub mask_grad: Option<TensorIr>,
1430    pub bias_grad: Option<TensorIr>,
1431}
1432
1433#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1434#[allow(missing_docs)]
1435pub struct Conv3dOpIr {
1436    pub x: TensorIr,
1437    pub weight: TensorIr,
1438    pub bias: Option<TensorIr>,
1439    pub options: Conv3dOptionsIr,
1440    pub out: TensorIr,
1441}
1442
1443#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1444#[allow(missing_docs)]
1445pub struct Conv3dXBackwardOpIr {
1446    pub x: TensorIr,
1447    pub weight: TensorIr,
1448    pub output_grad: TensorIr,
1449    pub options: Conv3dOptionsIr,
1450    pub out: TensorIr,
1451}
1452
1453#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1454#[allow(missing_docs)]
1455pub struct Conv3dWeightBackwardOpIr {
1456    pub x: TensorIr,
1457    pub weight: TensorIr,
1458    pub output_grad: TensorIr,
1459    pub options: Conv3dOptionsIr,
1460    pub out: TensorIr,
1461}
1462
1463#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1464#[allow(missing_docs)]
1465pub struct Conv3dBiasBackwardOpIr {
1466    pub x: TensorIr,
1467    pub bias: TensorIr,
1468    pub output_grad: TensorIr,
1469    pub out: TensorIr,
1470}
1471
1472#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1473#[allow(missing_docs)]
1474pub struct ConvTranspose1dOpIr {
1475    pub x: TensorIr,
1476    pub weight: TensorIr,
1477    pub bias: Option<TensorIr>,
1478    pub options: ConvTranspose1dOptionsIr,
1479    pub out: TensorIr,
1480}
1481
1482#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1483#[allow(missing_docs)]
1484pub struct ConvTranspose2dOpIr {
1485    pub x: TensorIr,
1486    pub weight: TensorIr,
1487    pub bias: Option<TensorIr>,
1488    pub options: ConvTranspose2dOptionsIr,
1489    pub out: TensorIr,
1490}
1491
1492#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1493#[allow(missing_docs)]
1494pub struct ConvTranspose3dOpIr {
1495    pub x: TensorIr,
1496    pub weight: TensorIr,
1497    pub bias: Option<TensorIr>,
1498    pub options: ConvTranspose3dOptionsIr,
1499    pub out: TensorIr,
1500}
1501
1502#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1503#[allow(missing_docs)]
1504pub struct Conv1dOptionsIr {
1505    pub stride: [usize; 1],
1506    pub padding: [(usize, usize); 1],
1507    pub dilation: [usize; 1],
1508    pub groups: usize,
1509}
1510
1511#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1512#[allow(missing_docs)]
1513pub struct Conv2dOptionsIr {
1514    pub stride: [usize; 2],
1515    pub padding: [(usize, usize); 2],
1516    pub dilation: [usize; 2],
1517    pub groups: usize,
1518}
1519
1520#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1521#[allow(missing_docs)]
1522pub struct DeformableConv2dOptionsIr {
1523    pub stride: [usize; 2],
1524    pub padding: [usize; 2],
1525    pub dilation: [usize; 2],
1526    pub weight_groups: usize,
1527    pub offset_groups: usize,
1528}
1529
1530#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1531#[allow(missing_docs)]
1532pub struct Conv3dOptionsIr {
1533    pub stride: [usize; 3],
1534    pub padding: [usize; 3],
1535    pub dilation: [usize; 3],
1536    pub groups: usize,
1537}
1538
1539#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1540#[allow(missing_docs)]
1541pub struct ConvTranspose1dOptionsIr {
1542    pub stride: [usize; 1],
1543    pub padding: [usize; 1],
1544    pub padding_out: [usize; 1],
1545    pub dilation: [usize; 1],
1546    pub groups: usize,
1547}
1548
1549#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1550#[allow(missing_docs)]
1551pub struct ConvTranspose2dOptionsIr {
1552    pub stride: [usize; 2],
1553    pub padding: [usize; 2],
1554    pub padding_out: [usize; 2],
1555    pub dilation: [usize; 2],
1556    pub groups: usize,
1557}
1558
1559#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1560#[allow(missing_docs)]
1561pub struct ConvTranspose3dOptionsIr {
1562    pub stride: [usize; 3],
1563    pub padding: [usize; 3],
1564    pub padding_out: [usize; 3],
1565    pub dilation: [usize; 3],
1566    pub groups: usize,
1567}
1568
1569/// Quantization parameters intermediate representation.
1570#[derive(Debug, Clone, Hash, PartialEq, Eq, Serialize, Deserialize)]
1571pub struct QuantizationParametersIr {
1572    /// The scaling factor, one per block or a single one for a per-tensor level.
1573    pub scales: TensorIr,
1574    /// The per-tensor scale that [`scales`](Self::scales) are expressed relative to, for a
1575    /// two-level scheme.
1576    pub global: Option<TensorIr>,
1577}
1578
1579#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1580#[allow(missing_docs)]
1581pub struct QuantizeOpIr {
1582    pub tensor: TensorIr,
1583    pub qparams: QuantizationParametersIr,
1584    pub scheme: QuantScheme,
1585    pub out: TensorIr,
1586}
1587
1588#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1589#[allow(missing_docs)]
1590pub struct DequantizeOpIr {
1591    pub input: TensorIr,
1592    pub out: TensorIr,
1593}
1594
1595impl From<ConvOptions<1>> for Conv1dOptionsIr {
1596    fn from(value: ConvOptions<1>) -> Self {
1597        Self {
1598            stride: value.stride,
1599            padding: value.padding,
1600            dilation: value.dilation,
1601            groups: value.groups,
1602        }
1603    }
1604}
1605
1606impl From<ConvOptions<2>> for Conv2dOptionsIr {
1607    fn from(value: ConvOptions<2>) -> Self {
1608        Self {
1609            stride: value.stride,
1610            padding: value.padding,
1611            dilation: value.dilation,
1612            groups: value.groups,
1613        }
1614    }
1615}
1616
1617impl From<ConvOptions<3>> for Conv3dOptionsIr {
1618    fn from(value: ConvOptions<3>) -> Self {
1619        Self {
1620            stride: value.stride,
1621            padding: value.symmetric_padding(),
1622            dilation: value.dilation,
1623            groups: value.groups,
1624        }
1625    }
1626}
1627
1628impl From<DeformConvOptions<2>> for DeformableConv2dOptionsIr {
1629    fn from(value: DeformConvOptions<2>) -> Self {
1630        Self {
1631            stride: value.stride,
1632            padding: value.padding,
1633            dilation: value.dilation,
1634            weight_groups: value.weight_groups,
1635            offset_groups: value.offset_groups,
1636        }
1637    }
1638}
1639
1640impl From<ConvTransposeOptions<1>> for ConvTranspose1dOptionsIr {
1641    fn from(value: ConvTransposeOptions<1>) -> Self {
1642        Self {
1643            stride: value.stride,
1644            padding: value.padding,
1645            padding_out: value.padding_out,
1646            dilation: value.dilation,
1647            groups: value.groups,
1648        }
1649    }
1650}
1651
1652impl From<ConvTransposeOptions<2>> for ConvTranspose2dOptionsIr {
1653    fn from(value: ConvTransposeOptions<2>) -> Self {
1654        Self {
1655            stride: value.stride,
1656            padding: value.padding,
1657            padding_out: value.padding_out,
1658            dilation: value.dilation,
1659            groups: value.groups,
1660        }
1661    }
1662}
1663
1664impl From<ConvTransposeOptions<3>> for ConvTranspose3dOptionsIr {
1665    fn from(value: ConvTransposeOptions<3>) -> Self {
1666        Self {
1667            stride: value.stride,
1668            padding: value.padding,
1669            padding_out: value.padding_out,
1670            dilation: value.dilation,
1671            groups: value.groups,
1672        }
1673    }
1674}
1675
1676impl From<Conv1dOptionsIr> for ConvOptions<1> {
1677    fn from(val: Conv1dOptionsIr) -> Self {
1678        ConvOptions {
1679            stride: val.stride,
1680            padding: val.padding,
1681            dilation: val.dilation,
1682            groups: val.groups,
1683        }
1684    }
1685}
1686
1687impl From<Conv2dOptionsIr> for ConvOptions<2> {
1688    fn from(val: Conv2dOptionsIr) -> Self {
1689        ConvOptions {
1690            stride: val.stride,
1691            padding: val.padding,
1692            dilation: val.dilation,
1693            groups: val.groups,
1694        }
1695    }
1696}
1697
1698impl From<Conv3dOptionsIr> for ConvOptions<3> {
1699    fn from(val: Conv3dOptionsIr) -> Self {
1700        ConvOptions {
1701            stride: val.stride,
1702            padding: val.padding.map(|padding| (padding, padding)),
1703            dilation: val.dilation,
1704            groups: val.groups,
1705        }
1706    }
1707}
1708
1709impl From<DeformableConv2dOptionsIr> for DeformConvOptions<2> {
1710    fn from(value: DeformableConv2dOptionsIr) -> Self {
1711        DeformConvOptions {
1712            stride: value.stride,
1713            padding: value.padding,
1714            dilation: value.dilation,
1715            weight_groups: value.weight_groups,
1716            offset_groups: value.offset_groups,
1717        }
1718    }
1719}
1720
1721impl From<ConvTranspose1dOptionsIr> for ConvTransposeOptions<1> {
1722    fn from(val: ConvTranspose1dOptionsIr) -> Self {
1723        ConvTransposeOptions {
1724            stride: val.stride,
1725            padding: val.padding,
1726            padding_out: val.padding_out,
1727            dilation: val.dilation,
1728            groups: val.groups,
1729        }
1730    }
1731}
1732
1733impl From<ConvTranspose2dOptionsIr> for ConvTransposeOptions<2> {
1734    fn from(val: ConvTranspose2dOptionsIr) -> Self {
1735        ConvTransposeOptions {
1736            stride: val.stride,
1737            padding: val.padding,
1738            padding_out: val.padding_out,
1739            dilation: val.dilation,
1740            groups: val.groups,
1741        }
1742    }
1743}
1744
1745impl From<ConvTranspose3dOptionsIr> for ConvTransposeOptions<3> {
1746    fn from(val: ConvTranspose3dOptionsIr) -> Self {
1747        ConvTransposeOptions {
1748            stride: val.stride,
1749            padding: val.padding,
1750            padding_out: val.padding_out,
1751            dilation: val.dilation,
1752            groups: val.groups,
1753        }
1754    }
1755}
1756
1757#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1758#[allow(missing_docs)]
1759pub struct AvgPool1dOpIr {
1760    pub x: TensorIr,
1761    pub kernel_size: usize,
1762    pub stride: usize,
1763    pub padding: usize,
1764    pub dilation: usize,
1765    pub count_include_pad: bool,
1766    pub ceil_mode: bool,
1767    pub out: TensorIr,
1768}
1769
1770#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1771#[allow(missing_docs)]
1772pub struct AvgPool2dOpIr {
1773    pub x: TensorIr,
1774    pub kernel_size: [usize; 2],
1775    pub stride: [usize; 2],
1776    pub padding: [usize; 2],
1777    pub dilation: [usize; 2],
1778    pub count_include_pad: bool,
1779    pub ceil_mode: bool,
1780    pub out: TensorIr,
1781}
1782
1783#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1784#[allow(missing_docs)]
1785pub struct AvgPool3dOpIr {
1786    pub x: TensorIr,
1787    pub kernel_size: [usize; 3],
1788    pub stride: [usize; 3],
1789    pub padding: [usize; 3],
1790    pub dilation: [usize; 3],
1791    pub count_include_pad: bool,
1792    pub ceil_mode: bool,
1793    pub out: TensorIr,
1794}
1795
1796#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1797#[allow(missing_docs)]
1798pub struct AvgPool1dBackwardOpIr {
1799    pub x: TensorIr,
1800    pub grad: TensorIr,
1801    pub kernel_size: usize,
1802    pub stride: usize,
1803    pub padding: usize,
1804    pub dilation: usize,
1805    pub count_include_pad: bool,
1806    pub ceil_mode: bool,
1807    pub out: TensorIr,
1808}
1809
1810#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1811#[allow(missing_docs)]
1812pub struct AvgPool2dBackwardOpIr {
1813    pub x: TensorIr,
1814    pub grad: TensorIr,
1815    pub kernel_size: [usize; 2],
1816    pub stride: [usize; 2],
1817    pub padding: [usize; 2],
1818    pub dilation: [usize; 2],
1819    pub count_include_pad: bool,
1820    pub ceil_mode: bool,
1821    pub out: TensorIr,
1822}
1823
1824#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1825#[allow(missing_docs)]
1826pub struct AvgPool3dBackwardOpIr {
1827    pub x: TensorIr,
1828    pub grad: TensorIr,
1829    pub kernel_size: [usize; 3],
1830    pub stride: [usize; 3],
1831    pub padding: [usize; 3],
1832    pub dilation: [usize; 3],
1833    pub count_include_pad: bool,
1834    pub ceil_mode: bool,
1835    pub out: TensorIr,
1836}
1837
1838#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1839#[allow(missing_docs)]
1840pub struct AdaptiveAvgPool1dOpIr {
1841    pub x: TensorIr,
1842    pub output_size: usize,
1843    pub out: TensorIr,
1844}
1845
1846#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1847#[allow(missing_docs)]
1848pub struct AdaptiveAvgPool2dOpIr {
1849    pub x: TensorIr,
1850    pub output_size: [usize; 2],
1851    pub out: TensorIr,
1852}
1853
1854#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1855#[allow(missing_docs)]
1856pub struct AdaptiveAvgPool1dBackwardOpIr {
1857    pub x: TensorIr,
1858    pub grad: TensorIr,
1859    pub out: TensorIr,
1860}
1861
1862#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1863#[allow(missing_docs)]
1864pub struct AdaptiveAvgPool2dBackwardOpIr {
1865    pub x: TensorIr,
1866    pub grad: TensorIr,
1867    pub out: TensorIr,
1868}
1869
1870#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1871#[allow(missing_docs)]
1872pub struct AdaptiveAvgPool3dOpIr {
1873    pub x: TensorIr,
1874    pub output_size: [usize; 3],
1875    pub out: TensorIr,
1876}
1877
1878#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1879#[allow(missing_docs)]
1880pub struct AdaptiveAvgPool3dBackwardOpIr {
1881    pub x: TensorIr,
1882    pub grad: TensorIr,
1883    pub out: TensorIr,
1884}
1885
1886#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1887#[allow(missing_docs)]
1888pub struct MaxPool1dOpIr {
1889    pub x: TensorIr,
1890    pub kernel_size: usize,
1891    pub stride: usize,
1892    pub padding: usize,
1893    pub dilation: usize,
1894    pub ceil_mode: bool,
1895    pub out: TensorIr,
1896}
1897
1898#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1899#[allow(missing_docs)]
1900pub struct MaxPool1dWithIndicesOpIr {
1901    pub x: TensorIr,
1902    pub kernel_size: usize,
1903    pub stride: usize,
1904    pub padding: usize,
1905    pub dilation: usize,
1906    pub ceil_mode: bool,
1907    pub out: TensorIr,
1908    pub out_indices: TensorIr,
1909}
1910
1911#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1912#[allow(missing_docs)]
1913pub struct MaxPool1dWithIndicesBackwardOpIr {
1914    pub x: TensorIr,
1915    pub grad: TensorIr,
1916    pub indices: TensorIr,
1917    pub kernel_size: usize,
1918    pub stride: usize,
1919    pub padding: usize,
1920    pub dilation: usize,
1921    pub ceil_mode: bool,
1922    pub out: TensorIr,
1923}
1924
1925#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1926#[allow(missing_docs)]
1927pub struct MaxPool2dOpIr {
1928    pub x: TensorIr,
1929    pub kernel_size: [usize; 2],
1930    pub stride: [usize; 2],
1931    pub padding: [usize; 2],
1932    pub dilation: [usize; 2],
1933    pub ceil_mode: bool,
1934    pub out: TensorIr,
1935}
1936
1937#[allow(missing_docs)]
1938#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1939pub struct MaxPool2dWithIndicesOpIr {
1940    pub x: TensorIr,
1941    pub kernel_size: [usize; 2],
1942    pub stride: [usize; 2],
1943    pub padding: [usize; 2],
1944    pub dilation: [usize; 2],
1945    pub ceil_mode: bool,
1946    pub out: TensorIr,
1947    pub out_indices: TensorIr,
1948}
1949
1950#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1951#[allow(missing_docs)]
1952pub struct MaxPool2dWithIndicesBackwardOpIr {
1953    pub x: TensorIr,
1954    pub grad: TensorIr,
1955    pub indices: TensorIr,
1956    pub kernel_size: [usize; 2],
1957    pub stride: [usize; 2],
1958    pub padding: [usize; 2],
1959    pub dilation: [usize; 2],
1960    pub ceil_mode: bool,
1961    pub out: TensorIr,
1962}
1963
1964#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1965#[allow(missing_docs)]
1966pub struct MaxPool3dOpIr {
1967    pub x: TensorIr,
1968    pub kernel_size: [usize; 3],
1969    pub stride: [usize; 3],
1970    pub padding: [usize; 3],
1971    pub dilation: [usize; 3],
1972    pub ceil_mode: bool,
1973    pub out: TensorIr,
1974}
1975
1976#[allow(missing_docs)]
1977#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1978pub struct MaxPool3dWithIndicesOpIr {
1979    pub x: TensorIr,
1980    pub kernel_size: [usize; 3],
1981    pub stride: [usize; 3],
1982    pub padding: [usize; 3],
1983    pub dilation: [usize; 3],
1984    pub ceil_mode: bool,
1985    pub out: TensorIr,
1986    pub out_indices: TensorIr,
1987}
1988
1989#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
1990#[allow(missing_docs)]
1991pub struct MaxPool3dWithIndicesBackwardOpIr {
1992    pub x: TensorIr,
1993    pub grad: TensorIr,
1994    pub indices: TensorIr,
1995    pub kernel_size: [usize; 3],
1996    pub stride: [usize; 3],
1997    pub padding: [usize; 3],
1998    pub dilation: [usize; 3],
1999    pub ceil_mode: bool,
2000    pub out: TensorIr,
2001}
2002
2003#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
2004#[allow(missing_docs)]
2005pub enum InterpolateModeIr {
2006    Nearest,
2007    NearestExact,
2008    Bilinear,
2009    Bicubic,
2010    Lanczos3,
2011}
2012
2013#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
2014#[allow(missing_docs)]
2015pub enum CoordinateTransformModeIr {
2016    HalfPixel,
2017    AlignCorners,
2018    Asymmetric,
2019    PytorchHalfPixel,
2020    HalfPixelSymmetric,
2021    TfCropAndResize {
2022        roi: [[f32; 2]; 2],
2023        extrapolation_value: f32,
2024    },
2025}
2026
2027fn canonicalize_f32_bits(val: f32) -> u32 {
2028    if val == 0.0 {
2029        0.0f32.to_bits()
2030    } else {
2031        val.to_bits()
2032    }
2033}
2034
2035impl Hash for CoordinateTransformModeIr {
2036    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
2037        core::mem::discriminant(self).hash(state);
2038        match self {
2039            Self::HalfPixel
2040            | Self::AlignCorners
2041            | Self::Asymmetric
2042            | Self::PytorchHalfPixel
2043            | Self::HalfPixelSymmetric => {}
2044            Self::TfCropAndResize {
2045                roi,
2046                extrapolation_value,
2047            } => {
2048                for row in roi {
2049                    for val in row {
2050                        canonicalize_f32_bits(*val).hash(state);
2051                    }
2052                }
2053                canonicalize_f32_bits(*extrapolation_value).hash(state);
2054            }
2055        }
2056    }
2057}
2058
2059#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, Serialize, Deserialize)]
2060#[allow(missing_docs)]
2061pub enum NearestRoundingModeIr {
2062    RoundPreferFloor,
2063    RoundPreferCeil,
2064    Floor,
2065    Ceil,
2066}
2067
2068#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
2069#[allow(missing_docs)]
2070pub struct CubicOptionsIr {
2071    pub a: f32,
2072    pub exclude_outside: bool,
2073}
2074
2075impl Hash for CubicOptionsIr {
2076    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
2077        canonicalize_f32_bits(self.a).hash(state);
2078        self.exclude_outside.hash(state);
2079    }
2080}
2081
2082#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
2083#[allow(missing_docs)]
2084pub struct InterpolateOptionsIr {
2085    pub mode: InterpolateModeIr,
2086    pub coordinate_transform: CoordinateTransformModeIr,
2087    pub nearest_mode: NearestRoundingModeIr,
2088    pub cubic_options: CubicOptionsIr,
2089    pub output_size: Option<[usize; 2]>,
2090    pub scale_factor: Option<[f32; 2]>,
2091}
2092
2093impl Hash for InterpolateOptionsIr {
2094    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
2095        self.mode.hash(state);
2096        self.coordinate_transform.hash(state);
2097        self.nearest_mode.hash(state);
2098        self.cubic_options.hash(state);
2099        self.output_size.hash(state);
2100        if let Some(scale) = self.scale_factor {
2101            for s in scale {
2102                canonicalize_f32_bits(s).hash(state);
2103            }
2104        }
2105    }
2106}
2107
2108#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
2109#[allow(missing_docs)]
2110pub struct InterpolateOpIr {
2111    pub x: TensorIr,
2112    pub output_size: [usize; 2],
2113    pub options: InterpolateOptionsIr,
2114    pub out: TensorIr,
2115}
2116
2117#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
2118#[allow(missing_docs)]
2119pub struct AttentionOptionsIr {
2120    pub scale: Option<ScalarIr>,
2121    pub softcap: Option<ScalarIr>,
2122    pub is_causal: bool,
2123    /// Missing in IR recorded before the field existed; that IR was bottom-right aligned.
2124    #[serde(default)]
2125    pub causal_alignment: CausalAlignment,
2126}
2127
2128impl From<AttentionOptionsIr> for AttentionModuleOptions {
2129    fn from(ir: AttentionOptionsIr) -> Self {
2130        AttentionModuleOptions {
2131            scale: ir.scale.map(|s| s.elem()),
2132            softcap: ir.softcap.map(|s| s.elem()),
2133            is_causal: ir.is_causal,
2134            causal_alignment: ir.causal_alignment,
2135        }
2136    }
2137}
2138
2139impl From<AttentionModuleOptions> for AttentionOptionsIr {
2140    fn from(ir: AttentionModuleOptions) -> Self {
2141        AttentionOptionsIr {
2142            scale: ir.scale.map(ScalarIr::Float),
2143            softcap: ir.softcap.map(ScalarIr::Float),
2144            is_causal: ir.is_causal,
2145            causal_alignment: ir.causal_alignment,
2146        }
2147    }
2148}
2149
2150#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
2151#[allow(missing_docs)]
2152pub struct AttentionOpIr {
2153    pub query: TensorIr,
2154    pub key: TensorIr,
2155    pub value: TensorIr,
2156    pub mask: Option<TensorIr>,
2157    pub attn_bias: Option<TensorIr>,
2158    pub options: AttentionOptionsIr,
2159    pub out: TensorIr,
2160}
2161
2162#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
2163#[allow(missing_docs)]
2164pub struct AttentionWithLogsumexpOpIr {
2165    pub query: TensorIr,
2166    pub key: TensorIr,
2167    pub value: TensorIr,
2168    pub options: AttentionOptionsIr,
2169    pub out: TensorIr,
2170    pub logsumexp: TensorIr,
2171}
2172
2173#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
2174#[allow(missing_docs)]
2175pub struct AttentionBackwardOpIr {
2176    pub query: TensorIr,
2177    pub key: TensorIr,
2178    pub value: TensorIr,
2179    pub out: TensorIr,
2180    pub logsumexp: TensorIr,
2181    pub out_grad: TensorIr,
2182    pub options: AttentionOptionsIr,
2183    pub query_grad: TensorIr,
2184    pub key_grad: TensorIr,
2185    pub value_grad: TensorIr,
2186}
2187
2188#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
2189#[allow(missing_docs)]
2190pub struct CtcLossOpIr {
2191    pub log_probs: TensorIr,
2192    pub targets: TensorIr,
2193    pub input_lengths: TensorIr,
2194    pub target_lengths: TensorIr,
2195    pub blank: usize,
2196    pub out: TensorIr,
2197}
2198
2199#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
2200#[allow(missing_docs)]
2201pub struct CtcLossBackwardOpIr {
2202    pub log_probs: TensorIr,
2203    pub targets: TensorIr,
2204    pub input_lengths: TensorIr,
2205    pub target_lengths: TensorIr,
2206    pub grad_loss: TensorIr,
2207    pub blank: usize,
2208    pub out: TensorIr,
2209}
2210
2211impl From<InterpolateModeIr> for InterpolateMode {
2212    fn from(val: InterpolateModeIr) -> Self {
2213        match val {
2214            InterpolateModeIr::Nearest => Self::Nearest,
2215            InterpolateModeIr::NearestExact => Self::NearestExact,
2216            InterpolateModeIr::Bilinear => Self::Bilinear,
2217            InterpolateModeIr::Bicubic => Self::Bicubic,
2218            InterpolateModeIr::Lanczos3 => Self::Lanczos3,
2219        }
2220    }
2221}
2222
2223impl From<CoordinateTransformModeIr> for CoordinateTransformMode {
2224    fn from(val: CoordinateTransformModeIr) -> Self {
2225        match val {
2226            CoordinateTransformModeIr::HalfPixel => Self::HalfPixel,
2227            CoordinateTransformModeIr::AlignCorners => Self::AlignCorners,
2228            CoordinateTransformModeIr::Asymmetric => Self::Asymmetric,
2229            CoordinateTransformModeIr::PytorchHalfPixel => Self::PytorchHalfPixel,
2230            CoordinateTransformModeIr::HalfPixelSymmetric => Self::HalfPixelSymmetric,
2231            CoordinateTransformModeIr::TfCropAndResize {
2232                roi,
2233                extrapolation_value,
2234            } => Self::TfCropAndResize {
2235                roi,
2236                extrapolation_value,
2237            },
2238        }
2239    }
2240}
2241
2242impl From<CoordinateTransformMode> for CoordinateTransformModeIr {
2243    fn from(val: CoordinateTransformMode) -> Self {
2244        match val {
2245            CoordinateTransformMode::HalfPixel => Self::HalfPixel,
2246            CoordinateTransformMode::AlignCorners => Self::AlignCorners,
2247            CoordinateTransformMode::Asymmetric => Self::Asymmetric,
2248            CoordinateTransformMode::PytorchHalfPixel => Self::PytorchHalfPixel,
2249            CoordinateTransformMode::HalfPixelSymmetric => Self::HalfPixelSymmetric,
2250            CoordinateTransformMode::TfCropAndResize {
2251                roi,
2252                extrapolation_value,
2253            } => Self::TfCropAndResize {
2254                roi,
2255                extrapolation_value,
2256            },
2257        }
2258    }
2259}
2260
2261impl From<NearestRoundingModeIr> for NearestRoundingMode {
2262    fn from(val: NearestRoundingModeIr) -> Self {
2263        match val {
2264            NearestRoundingModeIr::RoundPreferFloor => Self::RoundPreferFloor,
2265            NearestRoundingModeIr::RoundPreferCeil => Self::RoundPreferCeil,
2266            NearestRoundingModeIr::Floor => Self::Floor,
2267            NearestRoundingModeIr::Ceil => Self::Ceil,
2268        }
2269    }
2270}
2271
2272impl From<NearestRoundingMode> for NearestRoundingModeIr {
2273    fn from(val: NearestRoundingMode) -> Self {
2274        match val {
2275            NearestRoundingMode::RoundPreferFloor => Self::RoundPreferFloor,
2276            NearestRoundingMode::RoundPreferCeil => Self::RoundPreferCeil,
2277            NearestRoundingMode::Floor => Self::Floor,
2278            NearestRoundingMode::Ceil => Self::Ceil,
2279        }
2280    }
2281}
2282
2283impl From<CubicOptionsIr> for CubicOptions {
2284    fn from(val: CubicOptionsIr) -> Self {
2285        Self {
2286            a: val.a,
2287            exclude_outside: val.exclude_outside,
2288        }
2289    }
2290}
2291
2292impl From<CubicOptions> for CubicOptionsIr {
2293    fn from(val: CubicOptions) -> Self {
2294        Self {
2295            a: val.a,
2296            exclude_outside: val.exclude_outside,
2297        }
2298    }
2299}
2300
2301impl From<InterpolateOptionsIr> for InterpolateOptions {
2302    fn from(val: InterpolateOptionsIr) -> Self {
2303        Self {
2304            mode: val.mode.into(),
2305            coordinate_transform: val.coordinate_transform.into(),
2306            nearest_mode: val.nearest_mode.into(),
2307            cubic_options: val.cubic_options.into(),
2308            output_size: val.output_size,
2309            scale_factor: val.scale_factor,
2310        }
2311    }
2312}
2313
2314impl From<InterpolateMode> for InterpolateModeIr {
2315    fn from(val: InterpolateMode) -> Self {
2316        match val {
2317            InterpolateMode::Nearest => Self::Nearest,
2318            InterpolateMode::NearestExact => Self::NearestExact,
2319            InterpolateMode::Bilinear => Self::Bilinear,
2320            InterpolateMode::Bicubic => Self::Bicubic,
2321            InterpolateMode::Lanczos3 => Self::Lanczos3,
2322        }
2323    }
2324}
2325
2326impl From<InterpolateOptions> for InterpolateOptionsIr {
2327    fn from(val: InterpolateOptions) -> Self {
2328        Self {
2329            mode: val.mode.into(),
2330            coordinate_transform: val.coordinate_transform.into(),
2331            nearest_mode: val.nearest_mode.into(),
2332            cubic_options: val.cubic_options.into(),
2333            output_size: val.output_size,
2334            scale_factor: val.scale_factor,
2335        }
2336    }
2337}
2338
2339#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
2340#[allow(missing_docs)]
2341pub struct InterpolateBackwardOpIr {
2342    pub x: TensorIr,
2343    pub grad: TensorIr,
2344    pub output_size: [usize; 2],
2345    pub options: InterpolateOptionsIr,
2346    pub out: TensorIr,
2347}
2348
2349#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
2350#[allow(missing_docs)]
2351pub enum GridSamplePaddingModeIr {
2352    Zeros,
2353    Border,
2354    Reflection,
2355}
2356
2357#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
2358#[allow(missing_docs)]
2359pub struct GridSampleOptionsIr {
2360    pub mode: InterpolateModeIr,
2361    pub padding_mode: GridSamplePaddingModeIr,
2362    pub align_corners: bool,
2363}
2364
2365#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
2366#[allow(missing_docs)]
2367pub struct GridSample2dOpIr {
2368    pub tensor: TensorIr,
2369    pub grid: TensorIr,
2370    pub options: GridSampleOptionsIr,
2371    pub out: TensorIr,
2372}
2373
2374impl From<GridSamplePaddingModeIr> for GridSamplePaddingMode {
2375    fn from(val: GridSamplePaddingModeIr) -> Self {
2376        match val {
2377            GridSamplePaddingModeIr::Zeros => Self::Zeros,
2378            GridSamplePaddingModeIr::Border => Self::Border,
2379            GridSamplePaddingModeIr::Reflection => Self::Reflection,
2380        }
2381    }
2382}
2383
2384impl From<GridSamplePaddingMode> for GridSamplePaddingModeIr {
2385    fn from(val: GridSamplePaddingMode) -> Self {
2386        match val {
2387            GridSamplePaddingMode::Zeros => Self::Zeros,
2388            GridSamplePaddingMode::Border => Self::Border,
2389            GridSamplePaddingMode::Reflection => Self::Reflection,
2390        }
2391    }
2392}
2393
2394impl From<GridSampleOptionsIr> for GridSampleOptions {
2395    fn from(val: GridSampleOptionsIr) -> Self {
2396        Self {
2397            mode: val.mode.into(),
2398            padding_mode: val.padding_mode.into(),
2399            align_corners: val.align_corners,
2400        }
2401    }
2402}
2403
2404impl From<GridSampleOptions> for GridSampleOptionsIr {
2405    fn from(val: GridSampleOptions) -> Self {
2406        Self {
2407            mode: val.mode.into(),
2408            padding_mode: val.padding_mode.into(),
2409            align_corners: val.align_corners,
2410        }
2411    }
2412}
2413
2414impl OperationIr {
2415    /// Get all input [tensors](TensorIr) involved with the current operation.
2416    pub fn inputs(&self) -> impl Iterator<Item = &TensorIr> {
2417        match self {
2418            OperationIr::BaseFloat(repr) => repr.inputs(),
2419            OperationIr::BaseInt(repr) => repr.inputs(),
2420            OperationIr::BaseBool(repr) => repr.inputs(),
2421            OperationIr::NumericFloat(_dtype, repr) => repr.inputs(),
2422            OperationIr::NumericInt(_dtype, repr) => repr.inputs(),
2423            OperationIr::Bool(repr) => repr.inputs(),
2424            OperationIr::Int(repr) => repr.inputs(),
2425            OperationIr::Float(_dtype, repr) => repr.inputs(),
2426            OperationIr::Module(repr) => repr.inputs(),
2427            OperationIr::Init(repr) => repr.inputs(),
2428            OperationIr::Custom(repr) => repr.inputs(),
2429            OperationIr::Drop(repr) => Box::new([repr].into_iter()),
2430            OperationIr::Activation(repr) => repr.inputs(),
2431        }
2432    }
2433
2434    /// Get all output [tensors](TensorIr) involved with the current operation.
2435    pub fn outputs(&self) -> impl Iterator<Item = &TensorIr> {
2436        match self {
2437            OperationIr::BaseFloat(repr) => repr.outputs(),
2438            OperationIr::BaseInt(repr) => repr.outputs(),
2439            OperationIr::BaseBool(repr) => repr.outputs(),
2440            OperationIr::NumericFloat(_dtype, repr) => repr.outputs(),
2441            OperationIr::NumericInt(_dtype, repr) => repr.outputs(),
2442            OperationIr::Bool(repr) => repr.outputs(),
2443            OperationIr::Int(repr) => repr.outputs(),
2444            OperationIr::Float(_dtype, repr) => repr.outputs(),
2445            OperationIr::Module(repr) => repr.outputs(),
2446            OperationIr::Init(repr) => repr.outputs(),
2447            OperationIr::Custom(repr) => repr.outputs(),
2448            OperationIr::Drop(_repr) => Box::new([].into_iter()),
2449            OperationIr::Activation(repr) => repr.outputs(),
2450        }
2451    }
2452
2453    /// Get all [tensor](TensorIr) involved with the current operation.
2454    pub fn nodes(&self) -> Vec<&TensorIr> {
2455        self.inputs().chain(self.outputs()).collect()
2456    }
2457
2458    /// Set the given nodes that are [read write](super::TensorStatus::ReadWrite) to
2459    /// [read only](super::TensorStatus::ReadOnly) in the current operation.
2460    ///
2461    /// Returns the tensor that were updated with their original representation.
2462    pub fn mark_read_only(&mut self, nodes: &[TensorId]) -> Vec<TensorIr> {
2463        match self {
2464            OperationIr::BaseFloat(repr) => repr.mark_read_only(nodes),
2465            OperationIr::BaseInt(repr) => repr.mark_read_only(nodes),
2466            OperationIr::BaseBool(repr) => repr.mark_read_only(nodes),
2467            OperationIr::NumericFloat(_dtype, repr) => repr.mark_read_only(nodes),
2468            OperationIr::NumericInt(_dtype, repr) => repr.mark_read_only(nodes),
2469            OperationIr::Bool(repr) => repr.mark_read_only(nodes),
2470            OperationIr::Int(repr) => repr.mark_read_only(nodes),
2471            OperationIr::Float(_dtype, repr) => repr.mark_read_only(nodes),
2472            OperationIr::Module(repr) => repr.mark_read_only(nodes),
2473            OperationIr::Init(_) => Vec::new(),
2474            OperationIr::Drop(repr) => {
2475                let mut output = Vec::new();
2476                repr.mark_read_only(nodes, &mut output);
2477                output
2478            }
2479            OperationIr::Custom(repr) => {
2480                let mut output = Vec::new();
2481
2482                for input in repr.inputs.iter_mut() {
2483                    input.mark_read_only(nodes, &mut output);
2484                }
2485
2486                output
2487            }
2488            OperationIr::Activation(repr) => repr.mark_read_only(nodes),
2489        }
2490    }
2491
2492    /// Visit every mutable component of this operation in place.
2493    ///
2494    /// For each variant, visits its [tensors](TensorIr) (in the same order as
2495    /// [`inputs`](Self::inputs) then [`outputs`](Self::outputs) combined), then its
2496    /// [scalars](ScalarIr), then its slice ranges. Only `Slice`/`SliceAssign` (under the `Base*`
2497    /// variants) carry ranges; every other category visits no ranges.
2498    pub fn visit_mut(&mut self, v: &mut impl IrVisitorMut) {
2499        match self {
2500            OperationIr::BaseFloat(repr) => repr.visit_mut(v),
2501            OperationIr::BaseInt(repr) => repr.visit_mut(v),
2502            OperationIr::BaseBool(repr) => repr.visit_mut(v),
2503            OperationIr::NumericFloat(_dtype, repr) => repr.visit_mut(v),
2504            OperationIr::NumericInt(_dtype, repr) => repr.visit_mut(v),
2505            OperationIr::Bool(repr) => repr.visit_mut(v),
2506            OperationIr::Int(repr) => repr.visit_mut(v),
2507            OperationIr::Float(_dtype, repr) => repr.visit_mut(v),
2508            OperationIr::Module(repr) => repr.visit_mut(v),
2509            OperationIr::Init(repr) => repr.visit_mut(v),
2510            OperationIr::Custom(repr) => repr.visit_mut(v),
2511            OperationIr::Drop(repr) => v.visit_tensor_mut(repr),
2512            OperationIr::Activation(repr) => repr.visit_mut(v),
2513        }
2514    }
2515}
2516
2517impl BaseOperationIr {
2518    fn inputs(&self) -> Box<dyn Iterator<Item = &TensorIr> + '_> {
2519        match self {
2520            BaseOperationIr::Reshape(repr) => Box::new([&repr.input].into_iter()),
2521            BaseOperationIr::SwapDims(repr) => Box::new([&repr.input].into_iter()),
2522            BaseOperationIr::Permute(repr) => Box::new([&repr.input].into_iter()),
2523            BaseOperationIr::Expand(repr) => Box::new([&repr.input].into_iter()),
2524            BaseOperationIr::Flip(repr) => Box::new([&repr.input].into_iter()),
2525            BaseOperationIr::Slice(repr) => Box::new([&repr.tensor].into_iter()),
2526            BaseOperationIr::SliceAssign(repr) => Box::new([&repr.tensor, &repr.value].into_iter()),
2527            BaseOperationIr::Gather(repr) => Box::new([&repr.tensor, &repr.indices].into_iter()),
2528            BaseOperationIr::Scatter(repr) => {
2529                Box::new([&repr.tensor, &repr.indices, &repr.value].into_iter())
2530            }
2531            BaseOperationIr::ScatterNd(repr) => {
2532                Box::new([&repr.data, &repr.indices, &repr.values].into_iter())
2533            }
2534            BaseOperationIr::GatherNd(repr) => Box::new([&repr.data, &repr.indices].into_iter()),
2535            BaseOperationIr::Select(repr) => Box::new([&repr.tensor, &repr.indices].into_iter()),
2536            BaseOperationIr::SelectAssign(repr) => {
2537                Box::new([&repr.tensor, &repr.indices, &repr.value].into_iter())
2538            }
2539            BaseOperationIr::MaskWhere(repr) => {
2540                Box::new([&repr.tensor, &repr.mask, &repr.value].into_iter())
2541            }
2542            BaseOperationIr::MaskFill(repr) => Box::new([&repr.tensor, &repr.mask].into_iter()),
2543            BaseOperationIr::Equal(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()),
2544            BaseOperationIr::EqualElem(repr) => Box::new([&repr.lhs].into_iter()),
2545            BaseOperationIr::RepeatDim(repr) => Box::new([&repr.tensor].into_iter()),
2546            BaseOperationIr::Cat(repr) => Box::new(repr.tensors.iter()),
2547            BaseOperationIr::Cast(repr) => Box::new([&repr.input].into_iter()),
2548            BaseOperationIr::Unfold(repr) => Box::new([&repr.input].into_iter()),
2549            BaseOperationIr::Empty(_repr) => Box::new([].into_iter()),
2550            BaseOperationIr::Ones(_repr) => Box::new([].into_iter()),
2551            BaseOperationIr::Zeros(_repr) => Box::new([].into_iter()),
2552            BaseOperationIr::NotEqual(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()),
2553            BaseOperationIr::NotEqualElem(repr) => Box::new([&repr.lhs].into_iter()),
2554            BaseOperationIr::All(repr) => Box::new([&repr.input].into_iter()),
2555            BaseOperationIr::Any(repr) => Box::new([&repr.input].into_iter()),
2556            BaseOperationIr::AllDim(repr) => Box::new([&repr.input].into_iter()),
2557            BaseOperationIr::AnyDim(repr) => Box::new([&repr.input].into_iter()),
2558        }
2559    }
2560
2561    fn outputs(&self) -> Box<dyn Iterator<Item = &TensorIr> + '_> {
2562        match self {
2563            BaseOperationIr::Reshape(repr) => Box::new([&repr.out].into_iter()),
2564            BaseOperationIr::SwapDims(repr) => Box::new([&repr.out].into_iter()),
2565            BaseOperationIr::Permute(repr) => Box::new([&repr.out].into_iter()),
2566            BaseOperationIr::Expand(repr) => Box::new([&repr.out].into_iter()),
2567            BaseOperationIr::Flip(repr) => Box::new([&repr.out].into_iter()),
2568            BaseOperationIr::Slice(repr) => Box::new([&repr.out].into_iter()),
2569            BaseOperationIr::SliceAssign(repr) => Box::new([&repr.out].into_iter()),
2570            BaseOperationIr::Gather(repr) => Box::new([&repr.out].into_iter()),
2571            BaseOperationIr::Scatter(repr) => Box::new([&repr.out].into_iter()),
2572            BaseOperationIr::ScatterNd(repr) => Box::new([&repr.out].into_iter()),
2573            BaseOperationIr::GatherNd(repr) => Box::new([&repr.out].into_iter()),
2574            BaseOperationIr::Select(repr) => Box::new([&repr.out].into_iter()),
2575            BaseOperationIr::SelectAssign(repr) => Box::new([&repr.out].into_iter()),
2576            BaseOperationIr::MaskWhere(repr) => Box::new([&repr.out].into_iter()),
2577            BaseOperationIr::MaskFill(repr) => Box::new([&repr.out].into_iter()),
2578            BaseOperationIr::Equal(repr) => Box::new([&repr.out].into_iter()),
2579            BaseOperationIr::EqualElem(repr) => Box::new([&repr.out].into_iter()),
2580            BaseOperationIr::RepeatDim(repr) => Box::new([&repr.out].into_iter()),
2581            BaseOperationIr::Cat(repr) => Box::new([&repr.out].into_iter()),
2582            BaseOperationIr::Cast(repr) => Box::new([&repr.out].into_iter()),
2583            BaseOperationIr::Unfold(repr) => Box::new([&repr.out].into_iter()),
2584            BaseOperationIr::Empty(repr) => Box::new([&repr.out].into_iter()),
2585            BaseOperationIr::Ones(repr) => Box::new([&repr.out].into_iter()),
2586            BaseOperationIr::Zeros(repr) => Box::new([&repr.out].into_iter()),
2587            BaseOperationIr::NotEqual(repr) => Box::new([&repr.out].into_iter()),
2588            BaseOperationIr::NotEqualElem(repr) => Box::new([&repr.out].into_iter()),
2589            BaseOperationIr::All(repr) => Box::new([&repr.out].into_iter()),
2590            BaseOperationIr::Any(repr) => Box::new([&repr.out].into_iter()),
2591            BaseOperationIr::AllDim(repr) => Box::new([&repr.out].into_iter()),
2592            BaseOperationIr::AnyDim(repr) => Box::new([&repr.out].into_iter()),
2593        }
2594    }
2595
2596    fn mark_read_only(&mut self, nodes: &[TensorId]) -> Vec<TensorIr> {
2597        let mut output = Vec::new();
2598
2599        match self {
2600            BaseOperationIr::Reshape(repr) => {
2601                repr.input.mark_read_only(nodes, &mut output);
2602            }
2603            BaseOperationIr::SwapDims(repr) => {
2604                repr.input.mark_read_only(nodes, &mut output);
2605            }
2606            BaseOperationIr::Permute(repr) => {
2607                repr.input.mark_read_only(nodes, &mut output);
2608            }
2609
2610            BaseOperationIr::Expand(repr) => {
2611                repr.input.mark_read_only(nodes, &mut output);
2612            }
2613
2614            BaseOperationIr::Flip(repr) => {
2615                repr.input.mark_read_only(nodes, &mut output);
2616            }
2617            BaseOperationIr::Slice(repr) => {
2618                repr.tensor.mark_read_only(nodes, &mut output);
2619            }
2620            BaseOperationIr::SliceAssign(repr) => {
2621                repr.tensor.mark_read_only(nodes, &mut output);
2622                repr.value.mark_read_only(nodes, &mut output);
2623            }
2624            BaseOperationIr::Gather(repr) => {
2625                repr.tensor.mark_read_only(nodes, &mut output);
2626                repr.indices.mark_read_only(nodes, &mut output);
2627            }
2628            BaseOperationIr::Scatter(repr) => {
2629                repr.tensor.mark_read_only(nodes, &mut output);
2630                repr.indices.mark_read_only(nodes, &mut output);
2631                repr.value.mark_read_only(nodes, &mut output);
2632            }
2633            BaseOperationIr::ScatterNd(repr) => {
2634                repr.data.mark_read_only(nodes, &mut output);
2635                repr.indices.mark_read_only(nodes, &mut output);
2636                repr.values.mark_read_only(nodes, &mut output);
2637            }
2638            BaseOperationIr::GatherNd(repr) => {
2639                repr.data.mark_read_only(nodes, &mut output);
2640                repr.indices.mark_read_only(nodes, &mut output);
2641            }
2642            BaseOperationIr::Select(repr) => {
2643                repr.tensor.mark_read_only(nodes, &mut output);
2644                repr.indices.mark_read_only(nodes, &mut output);
2645            }
2646            BaseOperationIr::SelectAssign(repr) => {
2647                repr.tensor.mark_read_only(nodes, &mut output);
2648                repr.indices.mark_read_only(nodes, &mut output);
2649                repr.value.mark_read_only(nodes, &mut output);
2650            }
2651            BaseOperationIr::MaskWhere(repr) => {
2652                repr.tensor.mark_read_only(nodes, &mut output);
2653                repr.mask.mark_read_only(nodes, &mut output);
2654                repr.value.mark_read_only(nodes, &mut output);
2655            }
2656            BaseOperationIr::MaskFill(repr) => {
2657                repr.tensor.mark_read_only(nodes, &mut output);
2658                repr.mask.mark_read_only(nodes, &mut output);
2659            }
2660            BaseOperationIr::Equal(repr) => {
2661                repr.lhs.mark_read_only(nodes, &mut output);
2662                repr.rhs.mark_read_only(nodes, &mut output);
2663            }
2664            BaseOperationIr::EqualElem(repr) => {
2665                repr.lhs.mark_read_only(nodes, &mut output);
2666            }
2667            BaseOperationIr::RepeatDim(repr) => {
2668                repr.tensor.mark_read_only(nodes, &mut output);
2669            }
2670            BaseOperationIr::Cat(repr) => {
2671                for t in repr.tensors.iter_mut() {
2672                    t.mark_read_only(nodes, &mut output);
2673                }
2674            }
2675            BaseOperationIr::Cast(repr) => {
2676                repr.input.mark_read_only(nodes, &mut output);
2677            }
2678            BaseOperationIr::Unfold(repr) => {
2679                repr.input.mark_read_only(nodes, &mut output);
2680            }
2681            BaseOperationIr::Empty(_) => {}
2682            BaseOperationIr::Zeros(_) => {}
2683            BaseOperationIr::Ones(_) => {}
2684            BaseOperationIr::NotEqual(repr) => {
2685                repr.lhs.mark_read_only(nodes, &mut output);
2686                repr.rhs.mark_read_only(nodes, &mut output);
2687            }
2688            BaseOperationIr::NotEqualElem(repr) => {
2689                repr.lhs.mark_read_only(nodes, &mut output);
2690            }
2691            BaseOperationIr::All(repr) => {
2692                repr.input.mark_read_only(nodes, &mut output);
2693            }
2694            BaseOperationIr::Any(repr) => {
2695                repr.input.mark_read_only(nodes, &mut output);
2696            }
2697            BaseOperationIr::AllDim(repr) => {
2698                repr.input.mark_read_only(nodes, &mut output);
2699            }
2700            BaseOperationIr::AnyDim(repr) => {
2701                repr.input.mark_read_only(nodes, &mut output);
2702            }
2703        };
2704
2705        output
2706    }
2707
2708    fn visit_mut(&mut self, v: &mut impl IrVisitorMut) {
2709        match self {
2710            BaseOperationIr::Reshape(repr) => {
2711                v.visit_tensor_mut(&mut repr.input);
2712                v.visit_tensor_mut(&mut repr.out);
2713            }
2714            BaseOperationIr::SwapDims(repr) => {
2715                v.visit_tensor_mut(&mut repr.input);
2716                v.visit_tensor_mut(&mut repr.out);
2717            }
2718            BaseOperationIr::Permute(repr) => {
2719                v.visit_tensor_mut(&mut repr.input);
2720                v.visit_tensor_mut(&mut repr.out);
2721            }
2722            BaseOperationIr::Expand(repr) => {
2723                v.visit_tensor_mut(&mut repr.input);
2724                v.visit_tensor_mut(&mut repr.out);
2725            }
2726            BaseOperationIr::Flip(repr) => {
2727                v.visit_tensor_mut(&mut repr.input);
2728                v.visit_tensor_mut(&mut repr.out);
2729            }
2730            BaseOperationIr::Slice(repr) => {
2731                v.visit_tensor_mut(&mut repr.tensor);
2732                v.visit_tensor_mut(&mut repr.out);
2733                repr.ranges.iter_mut().for_each(|r| v.visit_range_mut(r));
2734            }
2735            BaseOperationIr::SliceAssign(repr) => {
2736                v.visit_tensor_mut(&mut repr.tensor);
2737                v.visit_tensor_mut(&mut repr.value);
2738                v.visit_tensor_mut(&mut repr.out);
2739                repr.ranges.iter_mut().for_each(|r| v.visit_range_mut(r));
2740            }
2741            BaseOperationIr::Gather(repr) => {
2742                v.visit_tensor_mut(&mut repr.tensor);
2743                v.visit_tensor_mut(&mut repr.indices);
2744                v.visit_tensor_mut(&mut repr.out);
2745            }
2746            BaseOperationIr::Scatter(repr) => {
2747                v.visit_tensor_mut(&mut repr.tensor);
2748                v.visit_tensor_mut(&mut repr.indices);
2749                v.visit_tensor_mut(&mut repr.value);
2750                v.visit_tensor_mut(&mut repr.out);
2751            }
2752            BaseOperationIr::ScatterNd(repr) => {
2753                v.visit_tensor_mut(&mut repr.data);
2754                v.visit_tensor_mut(&mut repr.indices);
2755                v.visit_tensor_mut(&mut repr.values);
2756                v.visit_tensor_mut(&mut repr.out);
2757            }
2758            BaseOperationIr::GatherNd(repr) => {
2759                v.visit_tensor_mut(&mut repr.data);
2760                v.visit_tensor_mut(&mut repr.indices);
2761                v.visit_tensor_mut(&mut repr.out);
2762            }
2763            BaseOperationIr::Select(repr) => {
2764                v.visit_tensor_mut(&mut repr.tensor);
2765                v.visit_tensor_mut(&mut repr.indices);
2766                v.visit_tensor_mut(&mut repr.out);
2767            }
2768            BaseOperationIr::SelectAssign(repr) => {
2769                v.visit_tensor_mut(&mut repr.tensor);
2770                v.visit_tensor_mut(&mut repr.indices);
2771                v.visit_tensor_mut(&mut repr.value);
2772                v.visit_tensor_mut(&mut repr.out);
2773            }
2774            BaseOperationIr::MaskWhere(repr) => {
2775                v.visit_tensor_mut(&mut repr.tensor);
2776                v.visit_tensor_mut(&mut repr.mask);
2777                v.visit_tensor_mut(&mut repr.value);
2778                v.visit_tensor_mut(&mut repr.out);
2779            }
2780            BaseOperationIr::MaskFill(repr) => {
2781                v.visit_tensor_mut(&mut repr.tensor);
2782                v.visit_tensor_mut(&mut repr.mask);
2783                v.visit_tensor_mut(&mut repr.out);
2784                v.visit_scalar_mut(&mut repr.value);
2785            }
2786            BaseOperationIr::Equal(repr) => {
2787                v.visit_tensor_mut(&mut repr.lhs);
2788                v.visit_tensor_mut(&mut repr.rhs);
2789                v.visit_tensor_mut(&mut repr.out);
2790            }
2791            BaseOperationIr::EqualElem(repr) => {
2792                v.visit_tensor_mut(&mut repr.lhs);
2793                v.visit_tensor_mut(&mut repr.out);
2794                v.visit_scalar_mut(&mut repr.rhs);
2795            }
2796            BaseOperationIr::RepeatDim(repr) => {
2797                v.visit_tensor_mut(&mut repr.tensor);
2798                v.visit_tensor_mut(&mut repr.out);
2799            }
2800            BaseOperationIr::Cat(repr) => {
2801                for t in repr.tensors.iter_mut() {
2802                    v.visit_tensor_mut(t);
2803                }
2804                v.visit_tensor_mut(&mut repr.out);
2805            }
2806            BaseOperationIr::Cast(repr) => {
2807                v.visit_tensor_mut(&mut repr.input);
2808                v.visit_tensor_mut(&mut repr.out);
2809            }
2810            BaseOperationIr::Unfold(repr) => {
2811                v.visit_tensor_mut(&mut repr.input);
2812                v.visit_tensor_mut(&mut repr.out);
2813            }
2814            BaseOperationIr::Empty(repr) => {
2815                v.visit_tensor_mut(&mut repr.out);
2816            }
2817            BaseOperationIr::Ones(repr) => {
2818                v.visit_tensor_mut(&mut repr.out);
2819            }
2820            BaseOperationIr::Zeros(repr) => {
2821                v.visit_tensor_mut(&mut repr.out);
2822            }
2823            BaseOperationIr::NotEqual(repr) => {
2824                v.visit_tensor_mut(&mut repr.lhs);
2825                v.visit_tensor_mut(&mut repr.rhs);
2826                v.visit_tensor_mut(&mut repr.out);
2827            }
2828            BaseOperationIr::NotEqualElem(repr) => {
2829                v.visit_tensor_mut(&mut repr.lhs);
2830                v.visit_tensor_mut(&mut repr.out);
2831                v.visit_scalar_mut(&mut repr.rhs);
2832            }
2833            BaseOperationIr::All(repr) => {
2834                v.visit_tensor_mut(&mut repr.input);
2835                v.visit_tensor_mut(&mut repr.out);
2836            }
2837            BaseOperationIr::Any(repr) => {
2838                v.visit_tensor_mut(&mut repr.input);
2839                v.visit_tensor_mut(&mut repr.out);
2840            }
2841            BaseOperationIr::AllDim(repr) => {
2842                v.visit_tensor_mut(&mut repr.input);
2843                v.visit_tensor_mut(&mut repr.out);
2844            }
2845            BaseOperationIr::AnyDim(repr) => {
2846                v.visit_tensor_mut(&mut repr.input);
2847                v.visit_tensor_mut(&mut repr.out);
2848            }
2849        }
2850    }
2851}
2852
2853impl NumericOperationIr {
2854    fn inputs(&self) -> Box<dyn Iterator<Item = &TensorIr> + '_> {
2855        match self {
2856            NumericOperationIr::Add(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()),
2857            NumericOperationIr::AddScalar(repr) => Box::new([&repr.lhs].into_iter()),
2858            NumericOperationIr::Sub(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()),
2859            NumericOperationIr::SubScalar(repr) => Box::new([&repr.lhs].into_iter()),
2860            NumericOperationIr::Mul(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()),
2861            NumericOperationIr::MulScalar(repr) => Box::new([&repr.lhs].into_iter()),
2862            NumericOperationIr::Div(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()),
2863            NumericOperationIr::DivScalar(repr) => Box::new([&repr.lhs].into_iter()),
2864            NumericOperationIr::Rem(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()),
2865            NumericOperationIr::RemScalar(repr) => Box::new([&repr.lhs].into_iter()),
2866            NumericOperationIr::GreaterElem(repr) => Box::new([&repr.lhs].into_iter()),
2867            NumericOperationIr::GreaterEqualElem(repr) => Box::new([&repr.lhs].into_iter()),
2868            NumericOperationIr::LowerElem(repr) => Box::new([&repr.lhs].into_iter()),
2869            NumericOperationIr::LowerEqualElem(repr) => Box::new([&repr.lhs].into_iter()),
2870            NumericOperationIr::Greater(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()),
2871            NumericOperationIr::GreaterEqual(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()),
2872            NumericOperationIr::Lower(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()),
2873            NumericOperationIr::LowerEqual(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()),
2874            NumericOperationIr::ArgMax(repr) => Box::new([&repr.input].into_iter()),
2875            NumericOperationIr::ArgTopK(repr) => Box::new([&repr.input].into_iter()),
2876            NumericOperationIr::TopK(repr) => Box::new([&repr.input].into_iter()),
2877            NumericOperationIr::ArgMin(repr) => Box::new([&repr.input].into_iter()),
2878            NumericOperationIr::Clamp(repr) => Box::new([&repr.tensor].into_iter()),
2879            NumericOperationIr::Abs(repr) => Box::new([&repr.input].into_iter()),
2880            NumericOperationIr::Full(_repr) => Box::new([].into_iter()),
2881            NumericOperationIr::MeanDim(repr) => Box::new([&repr.input].into_iter()),
2882            NumericOperationIr::Mean(repr) => Box::new([&repr.input].into_iter()),
2883            NumericOperationIr::Sum(repr) => Box::new([&repr.input].into_iter()),
2884            NumericOperationIr::SumDim(repr) => Box::new([&repr.input].into_iter()),
2885            NumericOperationIr::SumDims(repr) => Box::new([&repr.input].into_iter()),
2886            NumericOperationIr::Prod(repr) => Box::new([&repr.input].into_iter()),
2887            NumericOperationIr::ProdDim(repr) => Box::new([&repr.input].into_iter()),
2888            NumericOperationIr::Max(repr) => Box::new([&repr.input].into_iter()),
2889            NumericOperationIr::MaxDimWithIndices(repr) => Box::new([&repr.tensor].into_iter()),
2890            NumericOperationIr::TopKWithIndices(repr) => Box::new([&repr.tensor].into_iter()),
2891            NumericOperationIr::MinDimWithIndices(repr) => Box::new([&repr.tensor].into_iter()),
2892            NumericOperationIr::Min(repr) => Box::new([&repr.input].into_iter()),
2893            NumericOperationIr::MaxDim(repr) => Box::new([&repr.input].into_iter()),
2894            NumericOperationIr::MinDim(repr) => Box::new([&repr.input].into_iter()),
2895            NumericOperationIr::MaxAbs(repr) => Box::new([&repr.input].into_iter()),
2896            NumericOperationIr::MaxAbsDim(repr) => Box::new([&repr.input].into_iter()),
2897            NumericOperationIr::IntRandom(_repr) => Box::new([].into_iter()),
2898            NumericOperationIr::Powi(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()),
2899            NumericOperationIr::PowiScalar(repr) => Box::new([&repr.lhs].into_iter()),
2900            NumericOperationIr::CumMin(repr) => Box::new([&repr.input].into_iter()),
2901            NumericOperationIr::CumMax(repr) => Box::new([&repr.input].into_iter()),
2902            NumericOperationIr::CumProd(repr) => Box::new([&repr.input].into_iter()),
2903            NumericOperationIr::CumSum(repr) => Box::new([&repr.input].into_iter()),
2904            NumericOperationIr::Neg(repr) => Box::new([&repr.input].into_iter()),
2905            NumericOperationIr::Sign(repr) => Box::new([&repr.input].into_iter()),
2906            NumericOperationIr::ClampMin(repr) => Box::new([&repr.lhs].into_iter()),
2907            NumericOperationIr::ClampMax(repr) => Box::new([&repr.lhs].into_iter()),
2908            NumericOperationIr::Sort(repr) => Box::new([&repr.input].into_iter()),
2909            NumericOperationIr::SortWithIndices(repr) => Box::new([&repr.input].into_iter()),
2910            NumericOperationIr::ArgSort(repr) => Box::new([&repr.input].into_iter()),
2911            NumericOperationIr::Pad(repr) => Box::new([&repr.input].into_iter()),
2912        }
2913    }
2914
2915    fn outputs(&self) -> Box<dyn Iterator<Item = &TensorIr> + '_> {
2916        match self {
2917            NumericOperationIr::Add(repr) => Box::new([&repr.out].into_iter()),
2918            NumericOperationIr::AddScalar(repr) => Box::new([&repr.out].into_iter()),
2919            NumericOperationIr::Sub(repr) => Box::new([&repr.out].into_iter()),
2920            NumericOperationIr::SubScalar(repr) => Box::new([&repr.out].into_iter()),
2921            NumericOperationIr::Mul(repr) => Box::new([&repr.out].into_iter()),
2922            NumericOperationIr::MulScalar(repr) => Box::new([&repr.out].into_iter()),
2923            NumericOperationIr::Div(repr) => Box::new([&repr.out].into_iter()),
2924            NumericOperationIr::DivScalar(repr) => Box::new([&repr.out].into_iter()),
2925            NumericOperationIr::Rem(repr) => Box::new([&repr.out].into_iter()),
2926            NumericOperationIr::RemScalar(repr) => Box::new([&repr.out].into_iter()),
2927            NumericOperationIr::GreaterElem(repr) => Box::new([&repr.out].into_iter()),
2928            NumericOperationIr::GreaterEqualElem(repr) => Box::new([&repr.out].into_iter()),
2929            NumericOperationIr::LowerElem(repr) => Box::new([&repr.out].into_iter()),
2930            NumericOperationIr::LowerEqualElem(repr) => Box::new([&repr.out].into_iter()),
2931            NumericOperationIr::Greater(repr) => Box::new([&repr.out].into_iter()),
2932            NumericOperationIr::GreaterEqual(repr) => Box::new([&repr.out].into_iter()),
2933            NumericOperationIr::Lower(repr) => Box::new([&repr.out].into_iter()),
2934            NumericOperationIr::LowerEqual(repr) => Box::new([&repr.out].into_iter()),
2935            NumericOperationIr::ArgMax(repr) => Box::new([&repr.out].into_iter()),
2936            NumericOperationIr::ArgTopK(repr) => Box::new([&repr.out].into_iter()),
2937            NumericOperationIr::TopK(repr) => Box::new([&repr.out].into_iter()),
2938            NumericOperationIr::ArgMin(repr) => Box::new([&repr.out].into_iter()),
2939            NumericOperationIr::Clamp(repr) => Box::new([&repr.out].into_iter()),
2940            NumericOperationIr::Abs(repr) => Box::new([&repr.out].into_iter()),
2941            NumericOperationIr::Full(repr) => Box::new([&repr.out].into_iter()),
2942            NumericOperationIr::MeanDim(repr) => Box::new([&repr.out].into_iter()),
2943            NumericOperationIr::Mean(repr) => Box::new([&repr.out].into_iter()),
2944            NumericOperationIr::Sum(repr) => Box::new([&repr.out].into_iter()),
2945            NumericOperationIr::SumDim(repr) => Box::new([&repr.out].into_iter()),
2946            NumericOperationIr::SumDims(repr) => Box::new([&repr.out].into_iter()),
2947            NumericOperationIr::Prod(repr) => Box::new([&repr.out].into_iter()),
2948            NumericOperationIr::ProdDim(repr) => Box::new([&repr.out].into_iter()),
2949            NumericOperationIr::Max(repr) => Box::new([&repr.out].into_iter()),
2950            NumericOperationIr::MaxDimWithIndices(repr) => {
2951                Box::new([&repr.out, &repr.out_indices].into_iter())
2952            }
2953            NumericOperationIr::TopKWithIndices(repr) => {
2954                Box::new([&repr.out, &repr.out_indices].into_iter())
2955            }
2956            NumericOperationIr::MinDimWithIndices(repr) => {
2957                Box::new([&repr.out, &repr.out_indices].into_iter())
2958            }
2959            NumericOperationIr::Min(repr) => Box::new([&repr.out].into_iter()),
2960            NumericOperationIr::MaxDim(repr) => Box::new([&repr.out].into_iter()),
2961            NumericOperationIr::MinDim(repr) => Box::new([&repr.out].into_iter()),
2962            NumericOperationIr::MaxAbs(repr) => Box::new([&repr.out].into_iter()),
2963            NumericOperationIr::MaxAbsDim(repr) => Box::new([&repr.out].into_iter()),
2964            NumericOperationIr::IntRandom(repr) => Box::new([&repr.out].into_iter()),
2965            NumericOperationIr::Powi(repr) => Box::new([&repr.out].into_iter()),
2966            NumericOperationIr::PowiScalar(repr) => Box::new([&repr.out].into_iter()),
2967            NumericOperationIr::CumMin(repr) => Box::new([&repr.out].into_iter()),
2968            NumericOperationIr::CumMax(repr) => Box::new([&repr.out].into_iter()),
2969            NumericOperationIr::CumProd(repr) => Box::new([&repr.out].into_iter()),
2970            NumericOperationIr::CumSum(repr) => Box::new([&repr.out].into_iter()),
2971            NumericOperationIr::Neg(repr) => Box::new([&repr.out].into_iter()),
2972            NumericOperationIr::Sign(repr) => Box::new([&repr.out].into_iter()),
2973            NumericOperationIr::ClampMin(repr) => Box::new([&repr.out].into_iter()),
2974            NumericOperationIr::ClampMax(repr) => Box::new([&repr.out].into_iter()),
2975            NumericOperationIr::Sort(repr) => Box::new([&repr.out].into_iter()),
2976            NumericOperationIr::SortWithIndices(repr) => {
2977                Box::new([&repr.out, &repr.out_indices].into_iter())
2978            }
2979            NumericOperationIr::ArgSort(repr) => Box::new([&repr.out].into_iter()),
2980            NumericOperationIr::Pad(repr) => Box::new([&repr.out].into_iter()),
2981        }
2982    }
2983    fn mark_read_only(&mut self, nodes: &[TensorId]) -> Vec<TensorIr> {
2984        let mut output = Vec::new();
2985
2986        match self {
2987            NumericOperationIr::Add(repr) => {
2988                repr.lhs.mark_read_only(nodes, &mut output);
2989                repr.rhs.mark_read_only(nodes, &mut output);
2990            }
2991            NumericOperationIr::AddScalar(repr) => {
2992                repr.lhs.mark_read_only(nodes, &mut output);
2993            }
2994            NumericOperationIr::Sub(repr) => {
2995                repr.lhs.mark_read_only(nodes, &mut output);
2996                repr.rhs.mark_read_only(nodes, &mut output);
2997            }
2998            NumericOperationIr::SubScalar(repr) => {
2999                repr.lhs.mark_read_only(nodes, &mut output);
3000            }
3001            NumericOperationIr::Mul(repr) => {
3002                repr.lhs.mark_read_only(nodes, &mut output);
3003                repr.rhs.mark_read_only(nodes, &mut output);
3004            }
3005            NumericOperationIr::MulScalar(repr) => {
3006                repr.lhs.mark_read_only(nodes, &mut output);
3007            }
3008            NumericOperationIr::Div(repr) => {
3009                repr.lhs.mark_read_only(nodes, &mut output);
3010                repr.rhs.mark_read_only(nodes, &mut output);
3011            }
3012            NumericOperationIr::DivScalar(repr) => {
3013                repr.lhs.mark_read_only(nodes, &mut output);
3014            }
3015            NumericOperationIr::Rem(repr) => {
3016                repr.lhs.mark_read_only(nodes, &mut output);
3017                repr.rhs.mark_read_only(nodes, &mut output);
3018            }
3019            NumericOperationIr::RemScalar(repr) => {
3020                repr.lhs.mark_read_only(nodes, &mut output);
3021            }
3022            NumericOperationIr::GreaterElem(repr) => {
3023                repr.lhs.mark_read_only(nodes, &mut output);
3024            }
3025            NumericOperationIr::GreaterEqualElem(repr) => {
3026                repr.lhs.mark_read_only(nodes, &mut output);
3027            }
3028            NumericOperationIr::LowerElem(repr) => {
3029                repr.lhs.mark_read_only(nodes, &mut output);
3030            }
3031            NumericOperationIr::LowerEqualElem(repr) => {
3032                repr.lhs.mark_read_only(nodes, &mut output);
3033            }
3034            NumericOperationIr::Greater(repr) => {
3035                repr.lhs.mark_read_only(nodes, &mut output);
3036                repr.rhs.mark_read_only(nodes, &mut output);
3037            }
3038            NumericOperationIr::GreaterEqual(repr) => {
3039                repr.lhs.mark_read_only(nodes, &mut output);
3040                repr.rhs.mark_read_only(nodes, &mut output);
3041            }
3042            NumericOperationIr::Lower(repr) => {
3043                repr.lhs.mark_read_only(nodes, &mut output);
3044                repr.rhs.mark_read_only(nodes, &mut output);
3045            }
3046            NumericOperationIr::LowerEqual(repr) => {
3047                repr.lhs.mark_read_only(nodes, &mut output);
3048                repr.rhs.mark_read_only(nodes, &mut output);
3049            }
3050            NumericOperationIr::ArgMax(repr) => {
3051                repr.input.mark_read_only(nodes, &mut output);
3052            }
3053            NumericOperationIr::ArgTopK(repr) => {
3054                repr.input.mark_read_only(nodes, &mut output);
3055            }
3056            NumericOperationIr::TopK(repr) => {
3057                repr.input.mark_read_only(nodes, &mut output);
3058            }
3059            NumericOperationIr::ArgMin(repr) => {
3060                repr.input.mark_read_only(nodes, &mut output);
3061            }
3062            NumericOperationIr::Clamp(repr) => {
3063                repr.tensor.mark_read_only(nodes, &mut output);
3064            }
3065            NumericOperationIr::Abs(repr) => {
3066                repr.input.mark_read_only(nodes, &mut output);
3067            }
3068            NumericOperationIr::Full(_) => {}
3069            NumericOperationIr::MeanDim(repr) => {
3070                repr.input.mark_read_only(nodes, &mut output);
3071            }
3072            NumericOperationIr::Mean(repr) => {
3073                repr.input.mark_read_only(nodes, &mut output);
3074            }
3075            NumericOperationIr::Sum(repr) => {
3076                repr.input.mark_read_only(nodes, &mut output);
3077            }
3078            NumericOperationIr::SumDim(repr) => {
3079                repr.input.mark_read_only(nodes, &mut output);
3080            }
3081            NumericOperationIr::SumDims(repr) => {
3082                repr.input.mark_read_only(nodes, &mut output);
3083            }
3084            NumericOperationIr::Prod(repr) => {
3085                repr.input.mark_read_only(nodes, &mut output);
3086            }
3087            NumericOperationIr::ProdDim(repr) => {
3088                repr.input.mark_read_only(nodes, &mut output);
3089            }
3090            NumericOperationIr::Max(repr) => {
3091                repr.input.mark_read_only(nodes, &mut output);
3092            }
3093            NumericOperationIr::MaxDimWithIndices(repr) => {
3094                repr.tensor.mark_read_only(nodes, &mut output);
3095            }
3096            NumericOperationIr::TopKWithIndices(repr) => {
3097                repr.tensor.mark_read_only(nodes, &mut output);
3098            }
3099            NumericOperationIr::MinDimWithIndices(repr) => {
3100                repr.tensor.mark_read_only(nodes, &mut output);
3101            }
3102            NumericOperationIr::Min(repr) => {
3103                repr.input.mark_read_only(nodes, &mut output);
3104            }
3105            NumericOperationIr::MaxDim(repr) => {
3106                repr.input.mark_read_only(nodes, &mut output);
3107            }
3108            NumericOperationIr::MinDim(repr) => {
3109                repr.input.mark_read_only(nodes, &mut output);
3110            }
3111            NumericOperationIr::MaxAbs(repr) => {
3112                repr.input.mark_read_only(nodes, &mut output);
3113            }
3114            NumericOperationIr::MaxAbsDim(repr) => {
3115                repr.input.mark_read_only(nodes, &mut output);
3116            }
3117            NumericOperationIr::IntRandom(_) => {}
3118            NumericOperationIr::Powi(repr) => {
3119                repr.lhs.mark_read_only(nodes, &mut output);
3120                repr.rhs.mark_read_only(nodes, &mut output);
3121            }
3122            NumericOperationIr::PowiScalar(repr) => {
3123                repr.lhs.mark_read_only(nodes, &mut output);
3124            }
3125            NumericOperationIr::CumSum(repr) => {
3126                repr.input.mark_read_only(nodes, &mut output);
3127            }
3128            NumericOperationIr::CumProd(repr) => {
3129                repr.input.mark_read_only(nodes, &mut output);
3130            }
3131            NumericOperationIr::CumMin(repr) => {
3132                repr.input.mark_read_only(nodes, &mut output);
3133            }
3134            NumericOperationIr::CumMax(repr) => {
3135                repr.input.mark_read_only(nodes, &mut output);
3136            }
3137            NumericOperationIr::Neg(repr) => {
3138                repr.input.mark_read_only(nodes, &mut output);
3139            }
3140            NumericOperationIr::Sign(repr) => {
3141                repr.input.mark_read_only(nodes, &mut output);
3142            }
3143            NumericOperationIr::ClampMin(repr) => {
3144                repr.lhs.mark_read_only(nodes, &mut output);
3145            }
3146            NumericOperationIr::ClampMax(repr) => {
3147                repr.lhs.mark_read_only(nodes, &mut output);
3148            }
3149            NumericOperationIr::Sort(repr) => {
3150                repr.input.mark_read_only(nodes, &mut output);
3151            }
3152            NumericOperationIr::SortWithIndices(repr) => {
3153                repr.input.mark_read_only(nodes, &mut output);
3154            }
3155            NumericOperationIr::ArgSort(repr) => {
3156                repr.input.mark_read_only(nodes, &mut output);
3157            }
3158            NumericOperationIr::Pad(repr) => {
3159                repr.input.mark_read_only(nodes, &mut output);
3160            }
3161        };
3162
3163        output
3164    }
3165
3166    fn visit_mut(&mut self, v: &mut impl IrVisitorMut) {
3167        match self {
3168            NumericOperationIr::Add(repr) => {
3169                v.visit_tensor_mut(&mut repr.lhs);
3170                v.visit_tensor_mut(&mut repr.rhs);
3171                v.visit_tensor_mut(&mut repr.out);
3172            }
3173            NumericOperationIr::AddScalar(repr) => {
3174                v.visit_tensor_mut(&mut repr.lhs);
3175                v.visit_tensor_mut(&mut repr.out);
3176                v.visit_scalar_mut(&mut repr.rhs);
3177            }
3178            NumericOperationIr::Sub(repr) => {
3179                v.visit_tensor_mut(&mut repr.lhs);
3180                v.visit_tensor_mut(&mut repr.rhs);
3181                v.visit_tensor_mut(&mut repr.out);
3182            }
3183            NumericOperationIr::SubScalar(repr) => {
3184                v.visit_tensor_mut(&mut repr.lhs);
3185                v.visit_tensor_mut(&mut repr.out);
3186                v.visit_scalar_mut(&mut repr.rhs);
3187            }
3188            NumericOperationIr::Mul(repr) => {
3189                v.visit_tensor_mut(&mut repr.lhs);
3190                v.visit_tensor_mut(&mut repr.rhs);
3191                v.visit_tensor_mut(&mut repr.out);
3192            }
3193            NumericOperationIr::MulScalar(repr) => {
3194                v.visit_tensor_mut(&mut repr.lhs);
3195                v.visit_tensor_mut(&mut repr.out);
3196                v.visit_scalar_mut(&mut repr.rhs);
3197            }
3198            NumericOperationIr::Div(repr) => {
3199                v.visit_tensor_mut(&mut repr.lhs);
3200                v.visit_tensor_mut(&mut repr.rhs);
3201                v.visit_tensor_mut(&mut repr.out);
3202            }
3203            NumericOperationIr::DivScalar(repr) => {
3204                v.visit_tensor_mut(&mut repr.lhs);
3205                v.visit_tensor_mut(&mut repr.out);
3206                v.visit_scalar_mut(&mut repr.rhs);
3207            }
3208            NumericOperationIr::Rem(repr) => {
3209                v.visit_tensor_mut(&mut repr.lhs);
3210                v.visit_tensor_mut(&mut repr.rhs);
3211                v.visit_tensor_mut(&mut repr.out);
3212            }
3213            NumericOperationIr::RemScalar(repr) => {
3214                v.visit_tensor_mut(&mut repr.lhs);
3215                v.visit_tensor_mut(&mut repr.out);
3216                v.visit_scalar_mut(&mut repr.rhs);
3217            }
3218            NumericOperationIr::GreaterElem(repr) => {
3219                v.visit_tensor_mut(&mut repr.lhs);
3220                v.visit_tensor_mut(&mut repr.out);
3221                v.visit_scalar_mut(&mut repr.rhs);
3222            }
3223            NumericOperationIr::GreaterEqualElem(repr) => {
3224                v.visit_tensor_mut(&mut repr.lhs);
3225                v.visit_tensor_mut(&mut repr.out);
3226                v.visit_scalar_mut(&mut repr.rhs);
3227            }
3228            NumericOperationIr::LowerElem(repr) => {
3229                v.visit_tensor_mut(&mut repr.lhs);
3230                v.visit_tensor_mut(&mut repr.out);
3231                v.visit_scalar_mut(&mut repr.rhs);
3232            }
3233            NumericOperationIr::LowerEqualElem(repr) => {
3234                v.visit_tensor_mut(&mut repr.lhs);
3235                v.visit_tensor_mut(&mut repr.out);
3236                v.visit_scalar_mut(&mut repr.rhs);
3237            }
3238            NumericOperationIr::Greater(repr) => {
3239                v.visit_tensor_mut(&mut repr.lhs);
3240                v.visit_tensor_mut(&mut repr.rhs);
3241                v.visit_tensor_mut(&mut repr.out);
3242            }
3243            NumericOperationIr::GreaterEqual(repr) => {
3244                v.visit_tensor_mut(&mut repr.lhs);
3245                v.visit_tensor_mut(&mut repr.rhs);
3246                v.visit_tensor_mut(&mut repr.out);
3247            }
3248            NumericOperationIr::Lower(repr) => {
3249                v.visit_tensor_mut(&mut repr.lhs);
3250                v.visit_tensor_mut(&mut repr.rhs);
3251                v.visit_tensor_mut(&mut repr.out);
3252            }
3253            NumericOperationIr::LowerEqual(repr) => {
3254                v.visit_tensor_mut(&mut repr.lhs);
3255                v.visit_tensor_mut(&mut repr.rhs);
3256                v.visit_tensor_mut(&mut repr.out);
3257            }
3258            NumericOperationIr::ArgMax(repr) => {
3259                v.visit_tensor_mut(&mut repr.input);
3260                v.visit_tensor_mut(&mut repr.out);
3261            }
3262            NumericOperationIr::ArgTopK(repr) => {
3263                v.visit_tensor_mut(&mut repr.input);
3264                v.visit_tensor_mut(&mut repr.out);
3265            }
3266            NumericOperationIr::TopK(repr) => {
3267                v.visit_tensor_mut(&mut repr.input);
3268                v.visit_tensor_mut(&mut repr.out);
3269            }
3270            NumericOperationIr::ArgMin(repr) => {
3271                v.visit_tensor_mut(&mut repr.input);
3272                v.visit_tensor_mut(&mut repr.out);
3273            }
3274            NumericOperationIr::Clamp(repr) => {
3275                v.visit_tensor_mut(&mut repr.tensor);
3276                v.visit_tensor_mut(&mut repr.out);
3277                v.visit_scalar_mut(&mut repr.min);
3278                v.visit_scalar_mut(&mut repr.max);
3279            }
3280            NumericOperationIr::Abs(repr) => {
3281                v.visit_tensor_mut(&mut repr.input);
3282                v.visit_tensor_mut(&mut repr.out);
3283            }
3284            NumericOperationIr::Full(repr) => {
3285                v.visit_tensor_mut(&mut repr.out);
3286                v.visit_scalar_mut(&mut repr.value);
3287            }
3288            NumericOperationIr::MeanDim(repr) => {
3289                v.visit_tensor_mut(&mut repr.input);
3290                v.visit_tensor_mut(&mut repr.out);
3291            }
3292            NumericOperationIr::Mean(repr) => {
3293                v.visit_tensor_mut(&mut repr.input);
3294                v.visit_tensor_mut(&mut repr.out);
3295            }
3296            NumericOperationIr::Sum(repr) => {
3297                v.visit_tensor_mut(&mut repr.input);
3298                v.visit_tensor_mut(&mut repr.out);
3299            }
3300            NumericOperationIr::SumDim(repr) => {
3301                v.visit_tensor_mut(&mut repr.input);
3302                v.visit_tensor_mut(&mut repr.out);
3303            }
3304            NumericOperationIr::SumDims(repr) => {
3305                v.visit_tensor_mut(&mut repr.input);
3306                v.visit_tensor_mut(&mut repr.out);
3307            }
3308            NumericOperationIr::Prod(repr) => {
3309                v.visit_tensor_mut(&mut repr.input);
3310                v.visit_tensor_mut(&mut repr.out);
3311            }
3312            NumericOperationIr::ProdDim(repr) => {
3313                v.visit_tensor_mut(&mut repr.input);
3314                v.visit_tensor_mut(&mut repr.out);
3315            }
3316            NumericOperationIr::Max(repr) => {
3317                v.visit_tensor_mut(&mut repr.input);
3318                v.visit_tensor_mut(&mut repr.out);
3319            }
3320            NumericOperationIr::MaxDimWithIndices(repr) => {
3321                v.visit_tensor_mut(&mut repr.tensor);
3322                v.visit_tensor_mut(&mut repr.out);
3323                v.visit_tensor_mut(&mut repr.out_indices);
3324            }
3325            NumericOperationIr::TopKWithIndices(repr) => {
3326                v.visit_tensor_mut(&mut repr.tensor);
3327                v.visit_tensor_mut(&mut repr.out);
3328                v.visit_tensor_mut(&mut repr.out_indices);
3329            }
3330            NumericOperationIr::MinDimWithIndices(repr) => {
3331                v.visit_tensor_mut(&mut repr.tensor);
3332                v.visit_tensor_mut(&mut repr.out);
3333                v.visit_tensor_mut(&mut repr.out_indices);
3334            }
3335            NumericOperationIr::Min(repr) => {
3336                v.visit_tensor_mut(&mut repr.input);
3337                v.visit_tensor_mut(&mut repr.out);
3338            }
3339            NumericOperationIr::MaxDim(repr) => {
3340                v.visit_tensor_mut(&mut repr.input);
3341                v.visit_tensor_mut(&mut repr.out);
3342            }
3343            NumericOperationIr::MinDim(repr) => {
3344                v.visit_tensor_mut(&mut repr.input);
3345                v.visit_tensor_mut(&mut repr.out);
3346            }
3347            NumericOperationIr::MaxAbs(repr) => {
3348                v.visit_tensor_mut(&mut repr.input);
3349                v.visit_tensor_mut(&mut repr.out);
3350            }
3351            NumericOperationIr::MaxAbsDim(repr) => {
3352                v.visit_tensor_mut(&mut repr.input);
3353                v.visit_tensor_mut(&mut repr.out);
3354            }
3355            NumericOperationIr::IntRandom(repr) => {
3356                v.visit_tensor_mut(&mut repr.out);
3357            }
3358            NumericOperationIr::Powi(repr) => {
3359                v.visit_tensor_mut(&mut repr.lhs);
3360                v.visit_tensor_mut(&mut repr.rhs);
3361                v.visit_tensor_mut(&mut repr.out);
3362            }
3363            NumericOperationIr::PowiScalar(repr) => {
3364                v.visit_tensor_mut(&mut repr.lhs);
3365                v.visit_tensor_mut(&mut repr.out);
3366                v.visit_scalar_mut(&mut repr.rhs);
3367            }
3368            NumericOperationIr::CumMin(repr) => {
3369                v.visit_tensor_mut(&mut repr.input);
3370                v.visit_tensor_mut(&mut repr.out);
3371            }
3372            NumericOperationIr::CumMax(repr) => {
3373                v.visit_tensor_mut(&mut repr.input);
3374                v.visit_tensor_mut(&mut repr.out);
3375            }
3376            NumericOperationIr::CumProd(repr) => {
3377                v.visit_tensor_mut(&mut repr.input);
3378                v.visit_tensor_mut(&mut repr.out);
3379            }
3380            NumericOperationIr::CumSum(repr) => {
3381                v.visit_tensor_mut(&mut repr.input);
3382                v.visit_tensor_mut(&mut repr.out);
3383            }
3384            NumericOperationIr::Neg(repr) => {
3385                v.visit_tensor_mut(&mut repr.input);
3386                v.visit_tensor_mut(&mut repr.out);
3387            }
3388            NumericOperationIr::Sign(repr) => {
3389                v.visit_tensor_mut(&mut repr.input);
3390                v.visit_tensor_mut(&mut repr.out);
3391            }
3392            NumericOperationIr::ClampMin(repr) => {
3393                v.visit_tensor_mut(&mut repr.lhs);
3394                v.visit_tensor_mut(&mut repr.out);
3395                v.visit_scalar_mut(&mut repr.rhs);
3396            }
3397            NumericOperationIr::ClampMax(repr) => {
3398                v.visit_tensor_mut(&mut repr.lhs);
3399                v.visit_tensor_mut(&mut repr.out);
3400                v.visit_scalar_mut(&mut repr.rhs);
3401            }
3402            NumericOperationIr::Sort(repr) => {
3403                v.visit_tensor_mut(&mut repr.input);
3404                v.visit_tensor_mut(&mut repr.out);
3405            }
3406            NumericOperationIr::SortWithIndices(repr) => {
3407                v.visit_tensor_mut(&mut repr.input);
3408                v.visit_tensor_mut(&mut repr.out);
3409                v.visit_tensor_mut(&mut repr.out_indices);
3410            }
3411            NumericOperationIr::ArgSort(repr) => {
3412                v.visit_tensor_mut(&mut repr.input);
3413                v.visit_tensor_mut(&mut repr.out);
3414            }
3415            NumericOperationIr::Pad(repr) => {
3416                v.visit_tensor_mut(&mut repr.input);
3417                v.visit_tensor_mut(&mut repr.out);
3418                if let PadModeIr::Constant(value) = &mut repr.mode {
3419                    v.visit_scalar_mut(value);
3420                }
3421            }
3422        }
3423    }
3424}
3425
3426impl FloatOperationIr {
3427    fn inputs(&self) -> Box<dyn Iterator<Item = &TensorIr> + '_> {
3428        match self {
3429            FloatOperationIr::Matmul(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()),
3430            FloatOperationIr::Cross(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()),
3431            FloatOperationIr::Random(_repr) => Box::new([].into_iter()),
3432            FloatOperationIr::Exp(repr) => Box::new([&repr.input].into_iter()),
3433            FloatOperationIr::Log(repr) => Box::new([&repr.input].into_iter()),
3434            FloatOperationIr::Log1p(repr) => Box::new([&repr.input].into_iter()),
3435            FloatOperationIr::Erf(repr) => Box::new([&repr.input].into_iter()),
3436            FloatOperationIr::Recip(repr) => Box::new([&repr.input].into_iter()),
3437            FloatOperationIr::PowfScalar(repr) => Box::new([&repr.lhs].into_iter()),
3438            FloatOperationIr::Sqrt(repr) => Box::new([&repr.input].into_iter()),
3439            FloatOperationIr::Cos(repr) => Box::new([&repr.input].into_iter()),
3440            FloatOperationIr::Sin(repr) => Box::new([&repr.input].into_iter()),
3441            FloatOperationIr::Tanh(repr) => Box::new([&repr.input].into_iter()),
3442            FloatOperationIr::Round(repr) => Box::new([&repr.input].into_iter()),
3443            FloatOperationIr::Floor(repr) => Box::new([&repr.input].into_iter()),
3444            FloatOperationIr::Ceil(repr) => Box::new([&repr.input].into_iter()),
3445            FloatOperationIr::Trunc(repr) => Box::new([&repr.input].into_iter()),
3446            FloatOperationIr::IntoInt(repr) => Box::new([&repr.input].into_iter()),
3447            FloatOperationIr::Quantize(repr) => Box::new(
3448                [&repr.tensor, &repr.qparams.scales]
3449                    .into_iter()
3450                    .chain(repr.qparams.global.iter()),
3451            ),
3452            FloatOperationIr::Dequantize(repr) => Box::new([&repr.input].into_iter()),
3453            FloatOperationIr::IsNan(repr) => Box::new([&repr.input].into_iter()),
3454            FloatOperationIr::IsInf(repr) => Box::new([&repr.input].into_iter()),
3455            FloatOperationIr::GridSample2d(repr) => {
3456                Box::new([&repr.tensor, &repr.grid].into_iter())
3457            }
3458            FloatOperationIr::Tan(repr) => Box::new([&repr.input].into_iter()),
3459            FloatOperationIr::Cosh(repr) => Box::new([&repr.input].into_iter()),
3460            FloatOperationIr::Sinh(repr) => Box::new([&repr.input].into_iter()),
3461            FloatOperationIr::ArcCos(repr) => Box::new([&repr.input].into_iter()),
3462            FloatOperationIr::ArcCosh(repr) => Box::new([&repr.input].into_iter()),
3463            FloatOperationIr::ArcSin(repr) => Box::new([&repr.input].into_iter()),
3464            FloatOperationIr::ArcSinh(repr) => Box::new([&repr.input].into_iter()),
3465            FloatOperationIr::ArcTan(repr) => Box::new([&repr.input].into_iter()),
3466            FloatOperationIr::ArcTanh(repr) => Box::new([&repr.input].into_iter()),
3467            FloatOperationIr::ArcTan2(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()),
3468            FloatOperationIr::Powf(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()),
3469            FloatOperationIr::Hypot(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()),
3470            FloatOperationIr::Fmod(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()),
3471            FloatOperationIr::FmodScalar(repr) => Box::new([&repr.lhs].into_iter()),
3472        }
3473    }
3474    fn outputs(&self) -> Box<dyn Iterator<Item = &TensorIr> + '_> {
3475        match self {
3476            FloatOperationIr::Matmul(repr) => Box::new([&repr.out].into_iter()),
3477            FloatOperationIr::Cross(repr) => Box::new([&repr.out].into_iter()),
3478            FloatOperationIr::Random(repr) => Box::new([&repr.out].into_iter()),
3479            FloatOperationIr::Exp(repr) => Box::new([&repr.out].into_iter()),
3480            FloatOperationIr::Log(repr) => Box::new([&repr.out].into_iter()),
3481            FloatOperationIr::Log1p(repr) => Box::new([&repr.out].into_iter()),
3482            FloatOperationIr::Erf(repr) => Box::new([&repr.out].into_iter()),
3483            FloatOperationIr::Recip(repr) => Box::new([&repr.out].into_iter()),
3484            FloatOperationIr::PowfScalar(repr) => Box::new([&repr.out].into_iter()),
3485            FloatOperationIr::Sqrt(repr) => Box::new([&repr.out].into_iter()),
3486            FloatOperationIr::Cos(repr) => Box::new([&repr.out].into_iter()),
3487            FloatOperationIr::Sin(repr) => Box::new([&repr.out].into_iter()),
3488            FloatOperationIr::Tanh(repr) => Box::new([&repr.out].into_iter()),
3489            FloatOperationIr::Round(repr) => Box::new([&repr.out].into_iter()),
3490            FloatOperationIr::Floor(repr) => Box::new([&repr.out].into_iter()),
3491            FloatOperationIr::Ceil(repr) => Box::new([&repr.out].into_iter()),
3492            FloatOperationIr::Trunc(repr) => Box::new([&repr.out].into_iter()),
3493            FloatOperationIr::IntoInt(repr) => Box::new([&repr.out].into_iter()),
3494            FloatOperationIr::Quantize(repr) => Box::new([&repr.out].into_iter()),
3495            FloatOperationIr::Dequantize(repr) => Box::new([&repr.out].into_iter()),
3496            FloatOperationIr::IsNan(repr) => Box::new([&repr.out].into_iter()),
3497            FloatOperationIr::IsInf(repr) => Box::new([&repr.out].into_iter()),
3498            FloatOperationIr::GridSample2d(repr) => Box::new([&repr.out].into_iter()),
3499            FloatOperationIr::Tan(repr) => Box::new([&repr.out].into_iter()),
3500            FloatOperationIr::Cosh(repr) => Box::new([&repr.out].into_iter()),
3501            FloatOperationIr::Sinh(repr) => Box::new([&repr.out].into_iter()),
3502            FloatOperationIr::ArcCos(repr) => Box::new([&repr.out].into_iter()),
3503            FloatOperationIr::ArcCosh(repr) => Box::new([&repr.out].into_iter()),
3504            FloatOperationIr::ArcSin(repr) => Box::new([&repr.out].into_iter()),
3505            FloatOperationIr::ArcSinh(repr) => Box::new([&repr.out].into_iter()),
3506            FloatOperationIr::ArcTan(repr) => Box::new([&repr.out].into_iter()),
3507            FloatOperationIr::ArcTanh(repr) => Box::new([&repr.out].into_iter()),
3508            FloatOperationIr::ArcTan2(repr) => Box::new([&repr.out].into_iter()),
3509            FloatOperationIr::Powf(repr) => Box::new([&repr.out].into_iter()),
3510            FloatOperationIr::Hypot(repr) => Box::new([&repr.out].into_iter()),
3511            FloatOperationIr::Fmod(repr) => Box::new([&repr.out].into_iter()),
3512            FloatOperationIr::FmodScalar(repr) => Box::new([&repr.out].into_iter()),
3513        }
3514    }
3515
3516    fn mark_read_only(&mut self, nodes: &[TensorId]) -> Vec<TensorIr> {
3517        let mut output = Vec::new();
3518
3519        match self {
3520            FloatOperationIr::Matmul(repr) => {
3521                repr.lhs.mark_read_only(nodes, &mut output);
3522                repr.rhs.mark_read_only(nodes, &mut output);
3523            }
3524            FloatOperationIr::Cross(repr) => {
3525                repr.lhs.mark_read_only(nodes, &mut output);
3526                repr.rhs.mark_read_only(nodes, &mut output);
3527            }
3528            FloatOperationIr::Random(_) => {}
3529            FloatOperationIr::Exp(repr) => {
3530                repr.input.mark_read_only(nodes, &mut output);
3531            }
3532            FloatOperationIr::Log(repr) => {
3533                repr.input.mark_read_only(nodes, &mut output);
3534            }
3535            FloatOperationIr::Log1p(repr) => {
3536                repr.input.mark_read_only(nodes, &mut output);
3537            }
3538            FloatOperationIr::Erf(repr) => {
3539                repr.input.mark_read_only(nodes, &mut output);
3540            }
3541            FloatOperationIr::Recip(repr) => {
3542                repr.input.mark_read_only(nodes, &mut output);
3543            }
3544            FloatOperationIr::PowfScalar(repr) => {
3545                repr.lhs.mark_read_only(nodes, &mut output);
3546            }
3547            FloatOperationIr::Sqrt(repr) => {
3548                repr.input.mark_read_only(nodes, &mut output);
3549            }
3550            FloatOperationIr::Cos(repr) => {
3551                repr.input.mark_read_only(nodes, &mut output);
3552            }
3553            FloatOperationIr::Sin(repr) => {
3554                repr.input.mark_read_only(nodes, &mut output);
3555            }
3556            FloatOperationIr::Tanh(repr) => {
3557                repr.input.mark_read_only(nodes, &mut output);
3558            }
3559            FloatOperationIr::Round(repr) => {
3560                repr.input.mark_read_only(nodes, &mut output);
3561            }
3562            FloatOperationIr::Floor(repr) => {
3563                repr.input.mark_read_only(nodes, &mut output);
3564            }
3565            FloatOperationIr::Ceil(repr) => {
3566                repr.input.mark_read_only(nodes, &mut output);
3567            }
3568            FloatOperationIr::Trunc(repr) => {
3569                repr.input.mark_read_only(nodes, &mut output);
3570            }
3571            FloatOperationIr::Quantize(repr) => {
3572                repr.tensor.mark_read_only(nodes, &mut output);
3573                repr.qparams.scales.mark_read_only(nodes, &mut output);
3574                if let Some(global) = &mut repr.qparams.global {
3575                    global.mark_read_only(nodes, &mut output);
3576                }
3577            }
3578            FloatOperationIr::Dequantize(repr) => {
3579                repr.input.mark_read_only(nodes, &mut output);
3580            }
3581            FloatOperationIr::IntoInt(repr) => {
3582                repr.input.mark_read_only(nodes, &mut output);
3583            }
3584            FloatOperationIr::IsNan(repr) => {
3585                repr.input.mark_read_only(nodes, &mut output);
3586            }
3587            FloatOperationIr::IsInf(repr) => {
3588                repr.input.mark_read_only(nodes, &mut output);
3589            }
3590            FloatOperationIr::GridSample2d(repr) => {
3591                repr.tensor.mark_read_only(nodes, &mut output);
3592                repr.grid.mark_read_only(nodes, &mut output);
3593            }
3594            FloatOperationIr::Tan(repr) => repr.input.mark_read_only(nodes, &mut output),
3595            FloatOperationIr::Cosh(repr) => repr.input.mark_read_only(nodes, &mut output),
3596            FloatOperationIr::Sinh(repr) => repr.input.mark_read_only(nodes, &mut output),
3597            FloatOperationIr::ArcCos(repr) => repr.input.mark_read_only(nodes, &mut output),
3598            FloatOperationIr::ArcCosh(repr) => repr.input.mark_read_only(nodes, &mut output),
3599            FloatOperationIr::ArcSin(repr) => repr.input.mark_read_only(nodes, &mut output),
3600            FloatOperationIr::ArcSinh(repr) => repr.input.mark_read_only(nodes, &mut output),
3601            FloatOperationIr::ArcTan(repr) => repr.input.mark_read_only(nodes, &mut output),
3602            FloatOperationIr::ArcTanh(repr) => repr.input.mark_read_only(nodes, &mut output),
3603            FloatOperationIr::ArcTan2(repr) => {
3604                repr.lhs.mark_read_only(nodes, &mut output);
3605                repr.rhs.mark_read_only(nodes, &mut output);
3606            }
3607            FloatOperationIr::Powf(repr) => {
3608                repr.lhs.mark_read_only(nodes, &mut output);
3609                repr.rhs.mark_read_only(nodes, &mut output);
3610            }
3611            FloatOperationIr::Hypot(repr) => {
3612                repr.lhs.mark_read_only(nodes, &mut output);
3613                repr.rhs.mark_read_only(nodes, &mut output);
3614            }
3615            FloatOperationIr::Fmod(repr) => {
3616                repr.lhs.mark_read_only(nodes, &mut output);
3617                repr.rhs.mark_read_only(nodes, &mut output);
3618            }
3619            FloatOperationIr::FmodScalar(repr) => {
3620                repr.lhs.mark_read_only(nodes, &mut output);
3621            }
3622        };
3623
3624        output
3625    }
3626
3627    fn visit_mut(&mut self, v: &mut impl IrVisitorMut) {
3628        match self {
3629            FloatOperationIr::Matmul(repr) => {
3630                v.visit_tensor_mut(&mut repr.lhs);
3631                v.visit_tensor_mut(&mut repr.rhs);
3632                v.visit_tensor_mut(&mut repr.out);
3633            }
3634            FloatOperationIr::Cross(repr) => {
3635                v.visit_tensor_mut(&mut repr.lhs);
3636                v.visit_tensor_mut(&mut repr.rhs);
3637                v.visit_tensor_mut(&mut repr.out);
3638            }
3639            FloatOperationIr::Random(repr) => {
3640                v.visit_tensor_mut(&mut repr.out);
3641            }
3642            FloatOperationIr::Exp(repr) => {
3643                v.visit_tensor_mut(&mut repr.input);
3644                v.visit_tensor_mut(&mut repr.out);
3645            }
3646            FloatOperationIr::Log(repr) => {
3647                v.visit_tensor_mut(&mut repr.input);
3648                v.visit_tensor_mut(&mut repr.out);
3649            }
3650            FloatOperationIr::Log1p(repr) => {
3651                v.visit_tensor_mut(&mut repr.input);
3652                v.visit_tensor_mut(&mut repr.out);
3653            }
3654            FloatOperationIr::Erf(repr) => {
3655                v.visit_tensor_mut(&mut repr.input);
3656                v.visit_tensor_mut(&mut repr.out);
3657            }
3658            FloatOperationIr::Recip(repr) => {
3659                v.visit_tensor_mut(&mut repr.input);
3660                v.visit_tensor_mut(&mut repr.out);
3661            }
3662            FloatOperationIr::PowfScalar(repr) => {
3663                v.visit_tensor_mut(&mut repr.lhs);
3664                v.visit_tensor_mut(&mut repr.out);
3665                v.visit_scalar_mut(&mut repr.rhs);
3666            }
3667            FloatOperationIr::Sqrt(repr) => {
3668                v.visit_tensor_mut(&mut repr.input);
3669                v.visit_tensor_mut(&mut repr.out);
3670            }
3671            FloatOperationIr::Cos(repr) => {
3672                v.visit_tensor_mut(&mut repr.input);
3673                v.visit_tensor_mut(&mut repr.out);
3674            }
3675            FloatOperationIr::Sin(repr) => {
3676                v.visit_tensor_mut(&mut repr.input);
3677                v.visit_tensor_mut(&mut repr.out);
3678            }
3679            FloatOperationIr::Tanh(repr) => {
3680                v.visit_tensor_mut(&mut repr.input);
3681                v.visit_tensor_mut(&mut repr.out);
3682            }
3683            FloatOperationIr::Round(repr) => {
3684                v.visit_tensor_mut(&mut repr.input);
3685                v.visit_tensor_mut(&mut repr.out);
3686            }
3687            FloatOperationIr::Floor(repr) => {
3688                v.visit_tensor_mut(&mut repr.input);
3689                v.visit_tensor_mut(&mut repr.out);
3690            }
3691            FloatOperationIr::Ceil(repr) => {
3692                v.visit_tensor_mut(&mut repr.input);
3693                v.visit_tensor_mut(&mut repr.out);
3694            }
3695            FloatOperationIr::Trunc(repr) => {
3696                v.visit_tensor_mut(&mut repr.input);
3697                v.visit_tensor_mut(&mut repr.out);
3698            }
3699            FloatOperationIr::IntoInt(repr) => {
3700                v.visit_tensor_mut(&mut repr.input);
3701                v.visit_tensor_mut(&mut repr.out);
3702            }
3703            FloatOperationIr::Quantize(repr) => {
3704                v.visit_tensor_mut(&mut repr.tensor);
3705                v.visit_tensor_mut(&mut repr.qparams.scales);
3706                if let Some(global) = &mut repr.qparams.global {
3707                    v.visit_tensor_mut(global);
3708                }
3709                v.visit_tensor_mut(&mut repr.out);
3710            }
3711            FloatOperationIr::Dequantize(repr) => {
3712                v.visit_tensor_mut(&mut repr.input);
3713                v.visit_tensor_mut(&mut repr.out);
3714            }
3715            FloatOperationIr::IsNan(repr) => {
3716                v.visit_tensor_mut(&mut repr.input);
3717                v.visit_tensor_mut(&mut repr.out);
3718            }
3719            FloatOperationIr::IsInf(repr) => {
3720                v.visit_tensor_mut(&mut repr.input);
3721                v.visit_tensor_mut(&mut repr.out);
3722            }
3723            FloatOperationIr::GridSample2d(repr) => {
3724                v.visit_tensor_mut(&mut repr.tensor);
3725                v.visit_tensor_mut(&mut repr.grid);
3726                v.visit_tensor_mut(&mut repr.out);
3727            }
3728            FloatOperationIr::Tan(repr) => {
3729                v.visit_tensor_mut(&mut repr.input);
3730                v.visit_tensor_mut(&mut repr.out);
3731            }
3732            FloatOperationIr::Cosh(repr) => {
3733                v.visit_tensor_mut(&mut repr.input);
3734                v.visit_tensor_mut(&mut repr.out);
3735            }
3736            FloatOperationIr::Sinh(repr) => {
3737                v.visit_tensor_mut(&mut repr.input);
3738                v.visit_tensor_mut(&mut repr.out);
3739            }
3740            FloatOperationIr::ArcCos(repr) => {
3741                v.visit_tensor_mut(&mut repr.input);
3742                v.visit_tensor_mut(&mut repr.out);
3743            }
3744            FloatOperationIr::ArcCosh(repr) => {
3745                v.visit_tensor_mut(&mut repr.input);
3746                v.visit_tensor_mut(&mut repr.out);
3747            }
3748            FloatOperationIr::ArcSin(repr) => {
3749                v.visit_tensor_mut(&mut repr.input);
3750                v.visit_tensor_mut(&mut repr.out);
3751            }
3752            FloatOperationIr::ArcSinh(repr) => {
3753                v.visit_tensor_mut(&mut repr.input);
3754                v.visit_tensor_mut(&mut repr.out);
3755            }
3756            FloatOperationIr::ArcTan(repr) => {
3757                v.visit_tensor_mut(&mut repr.input);
3758                v.visit_tensor_mut(&mut repr.out);
3759            }
3760            FloatOperationIr::ArcTanh(repr) => {
3761                v.visit_tensor_mut(&mut repr.input);
3762                v.visit_tensor_mut(&mut repr.out);
3763            }
3764            FloatOperationIr::ArcTan2(repr) => {
3765                v.visit_tensor_mut(&mut repr.lhs);
3766                v.visit_tensor_mut(&mut repr.rhs);
3767                v.visit_tensor_mut(&mut repr.out);
3768            }
3769            FloatOperationIr::Powf(repr) => {
3770                v.visit_tensor_mut(&mut repr.lhs);
3771                v.visit_tensor_mut(&mut repr.rhs);
3772                v.visit_tensor_mut(&mut repr.out);
3773            }
3774            FloatOperationIr::Hypot(repr) => {
3775                v.visit_tensor_mut(&mut repr.lhs);
3776                v.visit_tensor_mut(&mut repr.rhs);
3777                v.visit_tensor_mut(&mut repr.out);
3778            }
3779            FloatOperationIr::Fmod(repr) => {
3780                v.visit_tensor_mut(&mut repr.lhs);
3781                v.visit_tensor_mut(&mut repr.rhs);
3782                v.visit_tensor_mut(&mut repr.out);
3783            }
3784            FloatOperationIr::FmodScalar(repr) => {
3785                v.visit_tensor_mut(&mut repr.lhs);
3786                v.visit_tensor_mut(&mut repr.out);
3787                v.visit_scalar_mut(&mut repr.rhs);
3788            }
3789        }
3790    }
3791}
3792
3793impl IntOperationIr {
3794    fn inputs(&self) -> Box<dyn Iterator<Item = &TensorIr> + '_> {
3795        match self {
3796            IntOperationIr::Matmul(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()),
3797            IntOperationIr::IntoFloat(repr) => Box::new([&repr.input].into_iter()),
3798            IntOperationIr::BitwiseAnd(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()),
3799            IntOperationIr::BitwiseAndScalar(repr) => Box::new([&repr.lhs].into_iter()),
3800            IntOperationIr::BitwiseOr(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()),
3801            IntOperationIr::BitwiseOrScalar(repr) => Box::new([&repr.lhs].into_iter()),
3802            IntOperationIr::BitwiseXor(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()),
3803            IntOperationIr::BitwiseXorScalar(repr) => Box::new([&repr.lhs].into_iter()),
3804            IntOperationIr::BitwiseNot(repr) => Box::new([&repr.input].into_iter()),
3805            IntOperationIr::BitwiseLeftShift(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()),
3806            IntOperationIr::BitwiseLeftShiftScalar(repr) => Box::new([&repr.lhs].into_iter()),
3807            IntOperationIr::BitwiseRightShift(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()),
3808            IntOperationIr::BitwiseRightShiftScalar(repr) => Box::new([&repr.lhs].into_iter()),
3809        }
3810    }
3811
3812    fn outputs(&self) -> Box<dyn Iterator<Item = &TensorIr> + '_> {
3813        match self {
3814            IntOperationIr::Matmul(repr) => Box::new([&repr.out].into_iter()),
3815            IntOperationIr::IntoFloat(repr) => Box::new([&repr.out].into_iter()),
3816            IntOperationIr::BitwiseAnd(repr) => Box::new([&repr.out].into_iter()),
3817            IntOperationIr::BitwiseAndScalar(repr) => Box::new([&repr.out].into_iter()),
3818            IntOperationIr::BitwiseOr(repr) => Box::new([&repr.out].into_iter()),
3819            IntOperationIr::BitwiseOrScalar(repr) => Box::new([&repr.out].into_iter()),
3820            IntOperationIr::BitwiseXor(repr) => Box::new([&repr.out].into_iter()),
3821            IntOperationIr::BitwiseXorScalar(repr) => Box::new([&repr.out].into_iter()),
3822            IntOperationIr::BitwiseNot(repr) => Box::new([&repr.out].into_iter()),
3823            IntOperationIr::BitwiseLeftShift(repr) => Box::new([&repr.out].into_iter()),
3824            IntOperationIr::BitwiseLeftShiftScalar(repr) => Box::new([&repr.out].into_iter()),
3825            IntOperationIr::BitwiseRightShift(repr) => Box::new([&repr.out].into_iter()),
3826            IntOperationIr::BitwiseRightShiftScalar(repr) => Box::new([&repr.out].into_iter()),
3827        }
3828    }
3829
3830    fn mark_read_only(&mut self, nodes: &[TensorId]) -> Vec<TensorIr> {
3831        let mut output = Vec::new();
3832
3833        match self {
3834            IntOperationIr::Matmul(repr) => {
3835                repr.lhs.mark_read_only(nodes, &mut output);
3836                repr.rhs.mark_read_only(nodes, &mut output);
3837            }
3838            IntOperationIr::IntoFloat(repr) => {
3839                repr.input.mark_read_only(nodes, &mut output);
3840            }
3841            IntOperationIr::BitwiseAnd(repr) => {
3842                repr.lhs.mark_read_only(nodes, &mut output);
3843                repr.rhs.mark_read_only(nodes, &mut output);
3844            }
3845            IntOperationIr::BitwiseAndScalar(repr) => {
3846                repr.lhs.mark_read_only(nodes, &mut output);
3847            }
3848            IntOperationIr::BitwiseOr(repr) => {
3849                repr.lhs.mark_read_only(nodes, &mut output);
3850                repr.rhs.mark_read_only(nodes, &mut output);
3851            }
3852            IntOperationIr::BitwiseOrScalar(repr) => {
3853                repr.lhs.mark_read_only(nodes, &mut output);
3854            }
3855            IntOperationIr::BitwiseXor(repr) => {
3856                repr.lhs.mark_read_only(nodes, &mut output);
3857                repr.rhs.mark_read_only(nodes, &mut output);
3858            }
3859            IntOperationIr::BitwiseXorScalar(repr) => {
3860                repr.lhs.mark_read_only(nodes, &mut output);
3861            }
3862            IntOperationIr::BitwiseNot(repr) => {
3863                repr.input.mark_read_only(nodes, &mut output);
3864            }
3865            IntOperationIr::BitwiseLeftShift(repr) => {
3866                repr.lhs.mark_read_only(nodes, &mut output);
3867                repr.rhs.mark_read_only(nodes, &mut output);
3868            }
3869            IntOperationIr::BitwiseLeftShiftScalar(repr) => {
3870                repr.lhs.mark_read_only(nodes, &mut output);
3871            }
3872            IntOperationIr::BitwiseRightShift(repr) => {
3873                repr.lhs.mark_read_only(nodes, &mut output);
3874                repr.rhs.mark_read_only(nodes, &mut output);
3875            }
3876            IntOperationIr::BitwiseRightShiftScalar(repr) => {
3877                repr.lhs.mark_read_only(nodes, &mut output);
3878            }
3879        };
3880
3881        output
3882    }
3883
3884    fn visit_mut(&mut self, v: &mut impl IrVisitorMut) {
3885        match self {
3886            IntOperationIr::Matmul(repr) => {
3887                v.visit_tensor_mut(&mut repr.lhs);
3888                v.visit_tensor_mut(&mut repr.rhs);
3889                v.visit_tensor_mut(&mut repr.out);
3890            }
3891            IntOperationIr::IntoFloat(repr) => {
3892                v.visit_tensor_mut(&mut repr.input);
3893                v.visit_tensor_mut(&mut repr.out);
3894            }
3895            IntOperationIr::BitwiseAnd(repr) => {
3896                v.visit_tensor_mut(&mut repr.lhs);
3897                v.visit_tensor_mut(&mut repr.rhs);
3898                v.visit_tensor_mut(&mut repr.out);
3899            }
3900            IntOperationIr::BitwiseAndScalar(repr) => {
3901                v.visit_tensor_mut(&mut repr.lhs);
3902                v.visit_tensor_mut(&mut repr.out);
3903                v.visit_scalar_mut(&mut repr.rhs);
3904            }
3905            IntOperationIr::BitwiseOr(repr) => {
3906                v.visit_tensor_mut(&mut repr.lhs);
3907                v.visit_tensor_mut(&mut repr.rhs);
3908                v.visit_tensor_mut(&mut repr.out);
3909            }
3910            IntOperationIr::BitwiseOrScalar(repr) => {
3911                v.visit_tensor_mut(&mut repr.lhs);
3912                v.visit_tensor_mut(&mut repr.out);
3913                v.visit_scalar_mut(&mut repr.rhs);
3914            }
3915            IntOperationIr::BitwiseXor(repr) => {
3916                v.visit_tensor_mut(&mut repr.lhs);
3917                v.visit_tensor_mut(&mut repr.rhs);
3918                v.visit_tensor_mut(&mut repr.out);
3919            }
3920            IntOperationIr::BitwiseXorScalar(repr) => {
3921                v.visit_tensor_mut(&mut repr.lhs);
3922                v.visit_tensor_mut(&mut repr.out);
3923                v.visit_scalar_mut(&mut repr.rhs);
3924            }
3925            IntOperationIr::BitwiseNot(repr) => {
3926                v.visit_tensor_mut(&mut repr.input);
3927                v.visit_tensor_mut(&mut repr.out);
3928            }
3929            IntOperationIr::BitwiseLeftShift(repr) => {
3930                v.visit_tensor_mut(&mut repr.lhs);
3931                v.visit_tensor_mut(&mut repr.rhs);
3932                v.visit_tensor_mut(&mut repr.out);
3933            }
3934            IntOperationIr::BitwiseLeftShiftScalar(repr) => {
3935                v.visit_tensor_mut(&mut repr.lhs);
3936                v.visit_tensor_mut(&mut repr.out);
3937                v.visit_scalar_mut(&mut repr.rhs);
3938            }
3939            IntOperationIr::BitwiseRightShift(repr) => {
3940                v.visit_tensor_mut(&mut repr.lhs);
3941                v.visit_tensor_mut(&mut repr.rhs);
3942                v.visit_tensor_mut(&mut repr.out);
3943            }
3944            IntOperationIr::BitwiseRightShiftScalar(repr) => {
3945                v.visit_tensor_mut(&mut repr.lhs);
3946                v.visit_tensor_mut(&mut repr.out);
3947                v.visit_scalar_mut(&mut repr.rhs);
3948            }
3949        }
3950    }
3951}
3952
3953impl BoolOperationIr {
3954    fn inputs(&self) -> Box<dyn Iterator<Item = &TensorIr> + '_> {
3955        match self {
3956            BoolOperationIr::IntoFloat(repr) => Box::new([&repr.input].into_iter()),
3957            BoolOperationIr::IntoInt(repr) => Box::new([&repr.input].into_iter()),
3958            BoolOperationIr::Not(repr) => Box::new([&repr.input].into_iter()),
3959            BoolOperationIr::And(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()),
3960            BoolOperationIr::Or(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()),
3961            BoolOperationIr::Xor(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()),
3962        }
3963    }
3964    fn outputs(&self) -> Box<dyn Iterator<Item = &TensorIr> + '_> {
3965        match self {
3966            BoolOperationIr::IntoFloat(repr) => Box::new([&repr.out].into_iter()),
3967            BoolOperationIr::IntoInt(repr) => Box::new([&repr.out].into_iter()),
3968            BoolOperationIr::Not(repr) => Box::new([&repr.out].into_iter()),
3969            BoolOperationIr::And(repr) => Box::new([&repr.out].into_iter()),
3970            BoolOperationIr::Or(repr) => Box::new([&repr.out].into_iter()),
3971            BoolOperationIr::Xor(repr) => Box::new([&repr.out].into_iter()),
3972        }
3973    }
3974    fn mark_read_only(&mut self, nodes: &[TensorId]) -> Vec<TensorIr> {
3975        let mut output = Vec::new();
3976
3977        match self {
3978            BoolOperationIr::IntoFloat(repr) => {
3979                repr.input.mark_read_only(nodes, &mut output);
3980            }
3981            BoolOperationIr::IntoInt(repr) => {
3982                repr.input.mark_read_only(nodes, &mut output);
3983            }
3984            BoolOperationIr::Not(repr) => {
3985                repr.input.mark_read_only(nodes, &mut output);
3986            }
3987            BoolOperationIr::And(repr) => {
3988                repr.lhs.mark_read_only(nodes, &mut output);
3989                repr.rhs.mark_read_only(nodes, &mut output);
3990            }
3991            BoolOperationIr::Or(repr) => {
3992                repr.lhs.mark_read_only(nodes, &mut output);
3993                repr.rhs.mark_read_only(nodes, &mut output);
3994            }
3995            BoolOperationIr::Xor(repr) => {
3996                repr.lhs.mark_read_only(nodes, &mut output);
3997                repr.rhs.mark_read_only(nodes, &mut output);
3998            }
3999        };
4000
4001        output
4002    }
4003
4004    fn visit_mut(&mut self, v: &mut impl IrVisitorMut) {
4005        match self {
4006            BoolOperationIr::IntoFloat(repr) => {
4007                v.visit_tensor_mut(&mut repr.input);
4008                v.visit_tensor_mut(&mut repr.out);
4009            }
4010            BoolOperationIr::IntoInt(repr) => {
4011                v.visit_tensor_mut(&mut repr.input);
4012                v.visit_tensor_mut(&mut repr.out);
4013            }
4014            BoolOperationIr::Not(repr) => {
4015                v.visit_tensor_mut(&mut repr.input);
4016                v.visit_tensor_mut(&mut repr.out);
4017            }
4018            BoolOperationIr::And(repr) => {
4019                v.visit_tensor_mut(&mut repr.lhs);
4020                v.visit_tensor_mut(&mut repr.rhs);
4021                v.visit_tensor_mut(&mut repr.out);
4022            }
4023            BoolOperationIr::Or(repr) => {
4024                v.visit_tensor_mut(&mut repr.lhs);
4025                v.visit_tensor_mut(&mut repr.rhs);
4026                v.visit_tensor_mut(&mut repr.out);
4027            }
4028            BoolOperationIr::Xor(repr) => {
4029                v.visit_tensor_mut(&mut repr.lhs);
4030                v.visit_tensor_mut(&mut repr.rhs);
4031                v.visit_tensor_mut(&mut repr.out);
4032            }
4033        }
4034    }
4035}
4036
4037impl ModuleOperationIr {
4038    fn inputs(&self) -> Box<dyn Iterator<Item = &TensorIr> + '_> {
4039        match self {
4040            ModuleOperationIr::BatchNorm(repr) => {
4041                Box::new([&repr.x, &repr.gamma, &repr.beta, &repr.mean, &repr.variance].into_iter())
4042            }
4043            ModuleOperationIr::Embedding(repr) => {
4044                Box::new([&repr.weights, &repr.indices].into_iter())
4045            }
4046            ModuleOperationIr::EmbeddingBackward(repr) => {
4047                Box::new([&repr.weights, &repr.out_grad, &repr.indices].into_iter())
4048            }
4049            ModuleOperationIr::Linear(repr) => {
4050                if let Some(bias) = &repr.bias {
4051                    Box::new([&repr.x, &repr.weight, bias].into_iter())
4052                } else {
4053                    Box::new([&repr.x, &repr.weight].into_iter())
4054                }
4055            }
4056            ModuleOperationIr::LinearXBackward(repr) => {
4057                Box::new([&repr.weight, &repr.output_grad].into_iter())
4058            }
4059            ModuleOperationIr::LinearWeightBackward(repr) => {
4060                Box::new([&repr.x, &repr.output_grad].into_iter())
4061            }
4062            ModuleOperationIr::LinearBiasBackward(repr) => {
4063                Box::new([&repr.output_grad].into_iter())
4064            }
4065            ModuleOperationIr::Conv1d(repr) => {
4066                if let Some(bias) = &repr.bias {
4067                    Box::new([&repr.x, &repr.weight, bias].into_iter())
4068                } else {
4069                    Box::new([&repr.x, &repr.weight].into_iter())
4070                }
4071            }
4072            ModuleOperationIr::Conv1dXBackward(repr) => {
4073                Box::new([&repr.x, &repr.weight, &repr.output_grad].into_iter())
4074            }
4075            ModuleOperationIr::Conv1dWeightBackward(repr) => {
4076                Box::new([&repr.x, &repr.weight, &repr.output_grad].into_iter())
4077            }
4078            ModuleOperationIr::Conv1dBiasBackward(repr) => {
4079                Box::new([&repr.x, &repr.bias, &repr.output_grad].into_iter())
4080            }
4081            ModuleOperationIr::Conv2d(repr) => {
4082                if let Some(bias) = &repr.bias {
4083                    Box::new([&repr.x, &repr.weight, bias].into_iter())
4084                } else {
4085                    Box::new([&repr.x, &repr.weight].into_iter())
4086                }
4087            }
4088            ModuleOperationIr::Conv2dXBackward(repr) => {
4089                Box::new([&repr.x, &repr.weight, &repr.output_grad].into_iter())
4090            }
4091            ModuleOperationIr::Conv2dWeightBackward(repr) => {
4092                Box::new([&repr.x, &repr.weight, &repr.output_grad].into_iter())
4093            }
4094            ModuleOperationIr::Conv2dBiasBackward(repr) => {
4095                Box::new([&repr.x, &repr.bias, &repr.output_grad].into_iter())
4096            }
4097            ModuleOperationIr::Conv3d(repr) => {
4098                if let Some(bias) = &repr.bias {
4099                    Box::new([&repr.x, &repr.weight, bias].into_iter())
4100                } else {
4101                    Box::new([&repr.x, &repr.weight].into_iter())
4102                }
4103            }
4104            ModuleOperationIr::Conv3dXBackward(repr) => {
4105                Box::new([&repr.x, &repr.weight, &repr.output_grad].into_iter())
4106            }
4107            ModuleOperationIr::Conv3dWeightBackward(repr) => {
4108                Box::new([&repr.x, &repr.weight, &repr.output_grad].into_iter())
4109            }
4110            ModuleOperationIr::Conv3dBiasBackward(repr) => {
4111                Box::new([&repr.x, &repr.bias, &repr.output_grad].into_iter())
4112            }
4113            ModuleOperationIr::DeformableConv2d(repr) => match (&repr.mask, &repr.bias) {
4114                (Some(mask), Some(bias)) => {
4115                    Box::new([&repr.x, &repr.offset, &repr.weight, mask, bias].into_iter())
4116                }
4117                (Some(mask), None) => {
4118                    Box::new([&repr.x, &repr.offset, &repr.weight, mask].into_iter())
4119                }
4120                (None, Some(bias)) => {
4121                    Box::new([&repr.x, &repr.offset, &repr.weight, bias].into_iter())
4122                }
4123                (None, None) => Box::new([&repr.x, &repr.offset, &repr.weight].into_iter()),
4124            },
4125            ModuleOperationIr::DeformableConv2dBackward(repr) => match (&repr.mask, &repr.bias) {
4126                (Some(mask), Some(bias)) => Box::new(
4127                    [
4128                        &repr.x,
4129                        &repr.offset,
4130                        &repr.weight,
4131                        &repr.out_grad,
4132                        mask,
4133                        bias,
4134                    ]
4135                    .into_iter(),
4136                ),
4137                (Some(mask), None) => Box::new(
4138                    [&repr.x, &repr.offset, &repr.weight, &repr.out_grad, mask].into_iter(),
4139                ),
4140                (None, Some(bias)) => Box::new(
4141                    [&repr.x, &repr.offset, &repr.weight, &repr.out_grad, bias].into_iter(),
4142                ),
4143                (None, None) => {
4144                    Box::new([&repr.x, &repr.offset, &repr.weight, &repr.out_grad].into_iter())
4145                }
4146            },
4147            ModuleOperationIr::ConvTranspose1d(repr) => {
4148                if let Some(bias) = &repr.bias {
4149                    Box::new([&repr.x, &repr.weight, bias].into_iter())
4150                } else {
4151                    Box::new([&repr.x, &repr.weight].into_iter())
4152                }
4153            }
4154            ModuleOperationIr::ConvTranspose2d(repr) => {
4155                if let Some(bias) = &repr.bias {
4156                    Box::new([&repr.x, &repr.weight, bias].into_iter())
4157                } else {
4158                    Box::new([&repr.x, &repr.weight].into_iter())
4159                }
4160            }
4161            ModuleOperationIr::ConvTranspose3d(repr) => {
4162                if let Some(bias) = &repr.bias {
4163                    Box::new([&repr.x, &repr.weight, bias].into_iter())
4164                } else {
4165                    Box::new([&repr.x, &repr.weight].into_iter())
4166                }
4167            }
4168            ModuleOperationIr::AvgPool1d(repr) => Box::new([&repr.x].into_iter()),
4169            ModuleOperationIr::AvgPool2d(repr) => Box::new([&repr.x].into_iter()),
4170            ModuleOperationIr::AvgPool3d(repr) => Box::new([&repr.x].into_iter()),
4171            ModuleOperationIr::AvgPool1dBackward(repr) => {
4172                Box::new([&repr.x, &repr.grad].into_iter())
4173            }
4174            ModuleOperationIr::AvgPool2dBackward(repr) => {
4175                Box::new([&repr.x, &repr.grad].into_iter())
4176            }
4177            ModuleOperationIr::AvgPool3dBackward(repr) => {
4178                Box::new([&repr.x, &repr.grad].into_iter())
4179            }
4180            ModuleOperationIr::AdaptiveAvgPool1d(repr) => Box::new([&repr.x].into_iter()),
4181            ModuleOperationIr::AdaptiveAvgPool2d(repr) => Box::new([&repr.x].into_iter()),
4182            ModuleOperationIr::AdaptiveAvgPool1dBackward(repr) => {
4183                Box::new([&repr.x, &repr.grad].into_iter())
4184            }
4185            ModuleOperationIr::AdaptiveAvgPool2dBackward(repr) => {
4186                Box::new([&repr.x, &repr.grad].into_iter())
4187            }
4188            ModuleOperationIr::AdaptiveAvgPool3d(repr) => Box::new([&repr.x].into_iter()),
4189            ModuleOperationIr::AdaptiveAvgPool3dBackward(repr) => {
4190                Box::new([&repr.x, &repr.grad].into_iter())
4191            }
4192            ModuleOperationIr::MaxPool1d(repr) => Box::new([&repr.x].into_iter()),
4193            ModuleOperationIr::MaxPool1dWithIndices(repr) => Box::new([&repr.x].into_iter()),
4194            ModuleOperationIr::MaxPool1dWithIndicesBackward(repr) => {
4195                Box::new([&repr.x, &repr.indices, &repr.grad].into_iter())
4196            }
4197            ModuleOperationIr::MaxPool2d(repr) => Box::new([&repr.x].into_iter()),
4198            ModuleOperationIr::MaxPool2dWithIndices(repr) => Box::new([&repr.x].into_iter()),
4199            ModuleOperationIr::MaxPool2dWithIndicesBackward(repr) => {
4200                Box::new([&repr.x, &repr.indices, &repr.grad].into_iter())
4201            }
4202            ModuleOperationIr::MaxPool3d(repr) => Box::new([&repr.x].into_iter()),
4203            ModuleOperationIr::MaxPool3dWithIndices(repr) => Box::new([&repr.x].into_iter()),
4204            ModuleOperationIr::MaxPool3dWithIndicesBackward(repr) => {
4205                Box::new([&repr.x, &repr.indices, &repr.grad].into_iter())
4206            }
4207            ModuleOperationIr::Interpolate(repr) => Box::new([&repr.x].into_iter()),
4208            ModuleOperationIr::InterpolateBackward(repr) => {
4209                Box::new([&repr.x, &repr.grad].into_iter())
4210            }
4211            ModuleOperationIr::Attention(repr) => {
4212                if let Some(mask) = &repr.mask {
4213                    if let Some(attn_bias) = &repr.attn_bias {
4214                        Box::new([&repr.query, &repr.key, &repr.value, mask, attn_bias].into_iter())
4215                    } else {
4216                        Box::new([&repr.query, &repr.key, &repr.value, mask].into_iter())
4217                    }
4218                } else if let Some(attn_bias) = &repr.attn_bias {
4219                    Box::new([&repr.query, &repr.key, &repr.value, attn_bias].into_iter())
4220                } else {
4221                    Box::new([&repr.query, &repr.key, &repr.value].into_iter())
4222                }
4223            }
4224            ModuleOperationIr::AttentionWithLogsumexp(repr) => {
4225                Box::new([&repr.query, &repr.key, &repr.value].into_iter())
4226            }
4227            ModuleOperationIr::AttentionBackward(repr) => Box::new(
4228                [
4229                    &repr.query,
4230                    &repr.key,
4231                    &repr.value,
4232                    &repr.out,
4233                    &repr.logsumexp,
4234                    &repr.out_grad,
4235                ]
4236                .into_iter(),
4237            ),
4238            ModuleOperationIr::CtcLoss(repr) => Box::new(
4239                [
4240                    &repr.log_probs,
4241                    &repr.targets,
4242                    &repr.input_lengths,
4243                    &repr.target_lengths,
4244                ]
4245                .into_iter(),
4246            ),
4247            ModuleOperationIr::CtcLossBackward(repr) => Box::new(
4248                [
4249                    &repr.log_probs,
4250                    &repr.targets,
4251                    &repr.input_lengths,
4252                    &repr.target_lengths,
4253                    &repr.grad_loss,
4254                ]
4255                .into_iter(),
4256            ),
4257            ModuleOperationIr::LayerNorm(repr) => match &repr.beta {
4258                Some(beta) => Box::new([&repr.input, &repr.gamma, beta].into_iter()),
4259                None => Box::new([&repr.input, &repr.gamma].into_iter()),
4260            },
4261            ModuleOperationIr::Unfold4d(repr) => Box::new([&repr.x].into_iter()),
4262            ModuleOperationIr::ConvTranspose1dWeightBackward(repr) => {
4263                Box::new([&repr.x, &repr.weight, &repr.output_grad].into_iter())
4264            }
4265            ModuleOperationIr::ConvTranspose1dBiasBackward(repr) => {
4266                Box::new([&repr.x, &repr.bias, &repr.output_grad].into_iter())
4267            }
4268            ModuleOperationIr::ConvTranspose2dWeightBackward(repr) => {
4269                Box::new([&repr.x, &repr.weight, &repr.output_grad].into_iter())
4270            }
4271            ModuleOperationIr::ConvTranspose2dBiasBackward(repr) => {
4272                Box::new([&repr.x, &repr.bias, &repr.output_grad].into_iter())
4273            }
4274            ModuleOperationIr::ConvTranspose3dWeightBackward(repr) => {
4275                Box::new([&repr.x, &repr.weight, &repr.output_grad].into_iter())
4276            }
4277            ModuleOperationIr::ConvTranspose3dBiasBackward(repr) => {
4278                Box::new([&repr.x, &repr.bias, &repr.output_grad].into_iter())
4279            }
4280        }
4281    }
4282    fn outputs(&self) -> Box<dyn Iterator<Item = &TensorIr> + '_> {
4283        match self {
4284            ModuleOperationIr::BatchNorm(repr) => Box::new([&repr.out].into_iter()),
4285            ModuleOperationIr::Embedding(repr) => Box::new([&repr.out].into_iter()),
4286            ModuleOperationIr::EmbeddingBackward(repr) => Box::new([&repr.out].into_iter()),
4287            ModuleOperationIr::Linear(repr) => Box::new([&repr.out].into_iter()),
4288            ModuleOperationIr::LinearXBackward(repr) => Box::new([&repr.out].into_iter()),
4289            ModuleOperationIr::LinearWeightBackward(repr) => Box::new([&repr.out].into_iter()),
4290            ModuleOperationIr::LinearBiasBackward(repr) => Box::new([&repr.out].into_iter()),
4291            ModuleOperationIr::Conv1d(repr) => Box::new([&repr.out].into_iter()),
4292            ModuleOperationIr::Conv1dXBackward(repr) => Box::new([&repr.out].into_iter()),
4293            ModuleOperationIr::Conv1dWeightBackward(repr) => Box::new([&repr.out].into_iter()),
4294            ModuleOperationIr::Conv1dBiasBackward(repr) => Box::new([&repr.out].into_iter()),
4295            ModuleOperationIr::Conv2d(repr) => Box::new([&repr.out].into_iter()),
4296            ModuleOperationIr::Conv2dXBackward(repr) => Box::new([&repr.out].into_iter()),
4297            ModuleOperationIr::Conv2dWeightBackward(repr) => Box::new([&repr.out].into_iter()),
4298            ModuleOperationIr::Conv2dBiasBackward(repr) => Box::new([&repr.out].into_iter()),
4299            ModuleOperationIr::Conv3d(repr) => Box::new([&repr.out].into_iter()),
4300            ModuleOperationIr::Conv3dXBackward(repr) => Box::new([&repr.out].into_iter()),
4301            ModuleOperationIr::Conv3dWeightBackward(repr) => Box::new([&repr.out].into_iter()),
4302            ModuleOperationIr::Conv3dBiasBackward(repr) => Box::new([&repr.out].into_iter()),
4303            ModuleOperationIr::DeformableConv2d(repr) => Box::new([&repr.out].into_iter()),
4304            ModuleOperationIr::DeformableConv2dBackward(repr) => {
4305                match (&repr.mask_grad, &repr.bias_grad) {
4306                    (Some(mask_grad), Some(bias_grad)) => Box::new(
4307                        [
4308                            &repr.input_grad,
4309                            &repr.offset_grad,
4310                            &repr.weight_grad,
4311                            mask_grad,
4312                            bias_grad,
4313                        ]
4314                        .into_iter(),
4315                    ),
4316                    (Some(mask_grad), None) => Box::new(
4317                        [
4318                            &repr.input_grad,
4319                            &repr.offset_grad,
4320                            &repr.weight_grad,
4321                            mask_grad,
4322                        ]
4323                        .into_iter(),
4324                    ),
4325                    (None, Some(bias_grad)) => Box::new(
4326                        [
4327                            &repr.input_grad,
4328                            &repr.offset_grad,
4329                            &repr.weight_grad,
4330                            bias_grad,
4331                        ]
4332                        .into_iter(),
4333                    ),
4334                    (None, None) => Box::new(
4335                        [&repr.input_grad, &repr.offset_grad, &repr.weight_grad].into_iter(),
4336                    ),
4337                }
4338            }
4339            ModuleOperationIr::ConvTranspose1d(repr) => Box::new([&repr.out].into_iter()),
4340            ModuleOperationIr::ConvTranspose2d(repr) => Box::new([&repr.out].into_iter()),
4341            ModuleOperationIr::ConvTranspose3d(repr) => Box::new([&repr.out].into_iter()),
4342            ModuleOperationIr::AvgPool1d(repr) => Box::new([&repr.out].into_iter()),
4343            ModuleOperationIr::AvgPool2d(repr) => Box::new([&repr.out].into_iter()),
4344            ModuleOperationIr::AvgPool3d(repr) => Box::new([&repr.out].into_iter()),
4345            ModuleOperationIr::AvgPool1dBackward(repr) => Box::new([&repr.out].into_iter()),
4346            ModuleOperationIr::AvgPool2dBackward(repr) => Box::new([&repr.out].into_iter()),
4347            ModuleOperationIr::AvgPool3dBackward(repr) => Box::new([&repr.out].into_iter()),
4348            ModuleOperationIr::AdaptiveAvgPool1d(repr) => Box::new([&repr.out].into_iter()),
4349            ModuleOperationIr::AdaptiveAvgPool2d(repr) => Box::new([&repr.out].into_iter()),
4350            ModuleOperationIr::AdaptiveAvgPool1dBackward(repr) => Box::new([&repr.out].into_iter()),
4351            ModuleOperationIr::AdaptiveAvgPool2dBackward(repr) => Box::new([&repr.out].into_iter()),
4352            ModuleOperationIr::AdaptiveAvgPool3d(repr) => Box::new([&repr.out].into_iter()),
4353            ModuleOperationIr::AdaptiveAvgPool3dBackward(repr) => Box::new([&repr.out].into_iter()),
4354            ModuleOperationIr::MaxPool1d(repr) => Box::new([&repr.out].into_iter()),
4355            ModuleOperationIr::MaxPool1dWithIndices(repr) => {
4356                Box::new([&repr.out, &repr.out_indices].into_iter())
4357            }
4358            ModuleOperationIr::MaxPool1dWithIndicesBackward(repr) => {
4359                Box::new([&repr.out].into_iter())
4360            }
4361            ModuleOperationIr::MaxPool2d(repr) => Box::new([&repr.out].into_iter()),
4362            ModuleOperationIr::MaxPool2dWithIndices(repr) => {
4363                Box::new([&repr.out, &repr.out_indices].into_iter())
4364            }
4365            ModuleOperationIr::MaxPool2dWithIndicesBackward(repr) => {
4366                Box::new([&repr.out].into_iter())
4367            }
4368            ModuleOperationIr::MaxPool3d(repr) => Box::new([&repr.out].into_iter()),
4369            ModuleOperationIr::MaxPool3dWithIndices(repr) => {
4370                Box::new([&repr.out, &repr.out_indices].into_iter())
4371            }
4372            ModuleOperationIr::MaxPool3dWithIndicesBackward(repr) => {
4373                Box::new([&repr.out].into_iter())
4374            }
4375            ModuleOperationIr::Interpolate(repr) => Box::new([&repr.out].into_iter()),
4376            ModuleOperationIr::InterpolateBackward(repr) => Box::new([&repr.out].into_iter()),
4377            ModuleOperationIr::Attention(repr) => Box::new([&repr.out].into_iter()),
4378            ModuleOperationIr::AttentionWithLogsumexp(repr) => {
4379                Box::new([&repr.out, &repr.logsumexp].into_iter())
4380            }
4381            ModuleOperationIr::AttentionBackward(repr) => {
4382                Box::new([&repr.query_grad, &repr.key_grad, &repr.value_grad].into_iter())
4383            }
4384            ModuleOperationIr::CtcLoss(repr) => Box::new([&repr.out].into_iter()),
4385            ModuleOperationIr::CtcLossBackward(repr) => Box::new([&repr.out].into_iter()),
4386            ModuleOperationIr::LayerNorm(repr) => Box::new([&repr.out].into_iter()),
4387            ModuleOperationIr::Unfold4d(repr) => Box::new([&repr.out].into_iter()),
4388            ModuleOperationIr::ConvTranspose1dWeightBackward(repr) => {
4389                Box::new([&repr.out].into_iter())
4390            }
4391            ModuleOperationIr::ConvTranspose1dBiasBackward(repr) => {
4392                Box::new([&repr.out].into_iter())
4393            }
4394            ModuleOperationIr::ConvTranspose2dWeightBackward(repr) => {
4395                Box::new([&repr.out].into_iter())
4396            }
4397            ModuleOperationIr::ConvTranspose2dBiasBackward(repr) => {
4398                Box::new([&repr.out].into_iter())
4399            }
4400            ModuleOperationIr::ConvTranspose3dWeightBackward(repr) => {
4401                Box::new([&repr.out].into_iter())
4402            }
4403            ModuleOperationIr::ConvTranspose3dBiasBackward(repr) => {
4404                Box::new([&repr.out].into_iter())
4405            }
4406        }
4407    }
4408
4409    fn mark_read_only(&mut self, nodes: &[TensorId]) -> Vec<TensorIr> {
4410        let mut output = Vec::new();
4411
4412        match self {
4413            ModuleOperationIr::BatchNorm(repr) => {
4414                repr.x.mark_read_only(nodes, &mut output);
4415                repr.gamma.mark_read_only(nodes, &mut output);
4416                repr.beta.mark_read_only(nodes, &mut output);
4417                repr.mean.mark_read_only(nodes, &mut output);
4418                repr.variance.mark_read_only(nodes, &mut output);
4419            }
4420            ModuleOperationIr::Embedding(repr) => {
4421                repr.weights.mark_read_only(nodes, &mut output);
4422                repr.indices.mark_read_only(nodes, &mut output);
4423            }
4424            ModuleOperationIr::EmbeddingBackward(repr) => {
4425                repr.weights.mark_read_only(nodes, &mut output);
4426                repr.out_grad.mark_read_only(nodes, &mut output);
4427                repr.indices.mark_read_only(nodes, &mut output);
4428            }
4429            ModuleOperationIr::Linear(repr) => {
4430                repr.x.mark_read_only(nodes, &mut output);
4431                repr.weight.mark_read_only(nodes, &mut output);
4432
4433                if let Some(bias) = &mut repr.bias {
4434                    bias.mark_read_only(nodes, &mut output);
4435                }
4436            }
4437            ModuleOperationIr::LinearXBackward(repr) => {
4438                repr.weight.mark_read_only(nodes, &mut output);
4439                repr.output_grad.mark_read_only(nodes, &mut output);
4440            }
4441            ModuleOperationIr::LinearWeightBackward(repr) => {
4442                repr.x.mark_read_only(nodes, &mut output);
4443                repr.output_grad.mark_read_only(nodes, &mut output);
4444            }
4445            ModuleOperationIr::LinearBiasBackward(repr) => {
4446                repr.output_grad.mark_read_only(nodes, &mut output);
4447            }
4448            ModuleOperationIr::Conv1d(repr) => {
4449                repr.x.mark_read_only(nodes, &mut output);
4450                repr.weight.mark_read_only(nodes, &mut output);
4451
4452                if let Some(bias) = &mut repr.bias {
4453                    bias.mark_read_only(nodes, &mut output);
4454                }
4455            }
4456            ModuleOperationIr::Conv1dXBackward(repr) => {
4457                repr.x.mark_read_only(nodes, &mut output);
4458                repr.weight.mark_read_only(nodes, &mut output);
4459                repr.output_grad.mark_read_only(nodes, &mut output);
4460            }
4461            ModuleOperationIr::Conv1dWeightBackward(repr) => {
4462                repr.x.mark_read_only(nodes, &mut output);
4463                repr.weight.mark_read_only(nodes, &mut output);
4464                repr.output_grad.mark_read_only(nodes, &mut output);
4465            }
4466            ModuleOperationIr::Conv1dBiasBackward(repr) => {
4467                repr.x.mark_read_only(nodes, &mut output);
4468                repr.bias.mark_read_only(nodes, &mut output);
4469                repr.output_grad.mark_read_only(nodes, &mut output);
4470            }
4471            ModuleOperationIr::Conv2d(repr) => {
4472                repr.x.mark_read_only(nodes, &mut output);
4473                repr.weight.mark_read_only(nodes, &mut output);
4474
4475                if let Some(bias) = &mut repr.bias {
4476                    bias.mark_read_only(nodes, &mut output);
4477                }
4478            }
4479            ModuleOperationIr::Conv2dXBackward(repr) => {
4480                repr.x.mark_read_only(nodes, &mut output);
4481                repr.weight.mark_read_only(nodes, &mut output);
4482                repr.output_grad.mark_read_only(nodes, &mut output);
4483            }
4484            ModuleOperationIr::Conv2dWeightBackward(repr) => {
4485                repr.x.mark_read_only(nodes, &mut output);
4486                repr.weight.mark_read_only(nodes, &mut output);
4487                repr.output_grad.mark_read_only(nodes, &mut output);
4488            }
4489            ModuleOperationIr::Conv2dBiasBackward(repr) => {
4490                repr.x.mark_read_only(nodes, &mut output);
4491                repr.bias.mark_read_only(nodes, &mut output);
4492                repr.output_grad.mark_read_only(nodes, &mut output);
4493            }
4494            ModuleOperationIr::Conv3d(repr) => {
4495                repr.x.mark_read_only(nodes, &mut output);
4496                repr.weight.mark_read_only(nodes, &mut output);
4497
4498                if let Some(bias) = &mut repr.bias {
4499                    bias.mark_read_only(nodes, &mut output);
4500                }
4501            }
4502            ModuleOperationIr::Conv3dXBackward(repr) => {
4503                repr.x.mark_read_only(nodes, &mut output);
4504                repr.weight.mark_read_only(nodes, &mut output);
4505                repr.output_grad.mark_read_only(nodes, &mut output);
4506            }
4507            ModuleOperationIr::Conv3dWeightBackward(repr) => {
4508                repr.x.mark_read_only(nodes, &mut output);
4509                repr.weight.mark_read_only(nodes, &mut output);
4510                repr.output_grad.mark_read_only(nodes, &mut output);
4511            }
4512            ModuleOperationIr::Conv3dBiasBackward(repr) => {
4513                repr.x.mark_read_only(nodes, &mut output);
4514                repr.bias.mark_read_only(nodes, &mut output);
4515                repr.output_grad.mark_read_only(nodes, &mut output);
4516            }
4517            ModuleOperationIr::DeformableConv2d(repr) => {
4518                repr.x.mark_read_only(nodes, &mut output);
4519                repr.weight.mark_read_only(nodes, &mut output);
4520                repr.offset.mark_read_only(nodes, &mut output);
4521
4522                match (&mut repr.mask, &mut repr.bias) {
4523                    (Some(mask), Some(bias)) => {
4524                        mask.mark_read_only(nodes, &mut output);
4525                        bias.mark_read_only(nodes, &mut output);
4526                    }
4527                    (Some(mask), None) => {
4528                        mask.mark_read_only(nodes, &mut output);
4529                    }
4530                    (None, Some(bias)) => {
4531                        bias.mark_read_only(nodes, &mut output);
4532                    }
4533                    (None, None) => {}
4534                };
4535            }
4536            ModuleOperationIr::DeformableConv2dBackward(repr) => {
4537                repr.x.mark_read_only(nodes, &mut output);
4538                repr.weight.mark_read_only(nodes, &mut output);
4539                repr.offset.mark_read_only(nodes, &mut output);
4540                repr.out_grad.mark_read_only(nodes, &mut output);
4541
4542                if let Some(mask) = repr.mask.as_mut() {
4543                    mask.mark_read_only(nodes, &mut output);
4544                }
4545                if let Some(bias) = repr.bias.as_mut() {
4546                    bias.mark_read_only(nodes, &mut output);
4547                }
4548            }
4549            ModuleOperationIr::ConvTranspose1d(repr) => {
4550                repr.x.mark_read_only(nodes, &mut output);
4551                repr.weight.mark_read_only(nodes, &mut output);
4552
4553                if let Some(bias) = &mut repr.bias {
4554                    bias.mark_read_only(nodes, &mut output);
4555                }
4556            }
4557            ModuleOperationIr::ConvTranspose2d(repr) => {
4558                repr.x.mark_read_only(nodes, &mut output);
4559                repr.weight.mark_read_only(nodes, &mut output);
4560
4561                if let Some(bias) = &mut repr.bias {
4562                    bias.mark_read_only(nodes, &mut output);
4563                }
4564            }
4565            ModuleOperationIr::ConvTranspose3d(repr) => {
4566                repr.x.mark_read_only(nodes, &mut output);
4567                repr.weight.mark_read_only(nodes, &mut output);
4568
4569                if let Some(bias) = &mut repr.bias {
4570                    bias.mark_read_only(nodes, &mut output);
4571                }
4572            }
4573            ModuleOperationIr::AvgPool1d(repr) => {
4574                repr.x.mark_read_only(nodes, &mut output);
4575            }
4576            ModuleOperationIr::AvgPool2d(repr) => {
4577                repr.x.mark_read_only(nodes, &mut output);
4578            }
4579            ModuleOperationIr::AvgPool1dBackward(repr) => {
4580                repr.x.mark_read_only(nodes, &mut output);
4581                repr.grad.mark_read_only(nodes, &mut output);
4582            }
4583            ModuleOperationIr::AvgPool2dBackward(repr) => {
4584                repr.x.mark_read_only(nodes, &mut output);
4585                repr.grad.mark_read_only(nodes, &mut output);
4586            }
4587            ModuleOperationIr::AdaptiveAvgPool1d(repr) => {
4588                repr.x.mark_read_only(nodes, &mut output);
4589            }
4590            ModuleOperationIr::AdaptiveAvgPool2d(repr) => {
4591                repr.x.mark_read_only(nodes, &mut output);
4592            }
4593            ModuleOperationIr::AdaptiveAvgPool1dBackward(repr) => {
4594                repr.x.mark_read_only(nodes, &mut output);
4595                repr.grad.mark_read_only(nodes, &mut output);
4596            }
4597            ModuleOperationIr::AdaptiveAvgPool2dBackward(repr) => {
4598                repr.x.mark_read_only(nodes, &mut output);
4599                repr.grad.mark_read_only(nodes, &mut output);
4600            }
4601            ModuleOperationIr::AdaptiveAvgPool3d(repr) => {
4602                repr.x.mark_read_only(nodes, &mut output);
4603            }
4604            ModuleOperationIr::AdaptiveAvgPool3dBackward(repr) => {
4605                repr.x.mark_read_only(nodes, &mut output);
4606                repr.grad.mark_read_only(nodes, &mut output);
4607            }
4608            ModuleOperationIr::MaxPool1d(repr) => {
4609                repr.x.mark_read_only(nodes, &mut output);
4610            }
4611            ModuleOperationIr::MaxPool1dWithIndices(repr) => {
4612                repr.x.mark_read_only(nodes, &mut output);
4613            }
4614            ModuleOperationIr::MaxPool1dWithIndicesBackward(repr) => {
4615                repr.x.mark_read_only(nodes, &mut output);
4616                repr.grad.mark_read_only(nodes, &mut output);
4617            }
4618            ModuleOperationIr::MaxPool2d(repr) => {
4619                repr.x.mark_read_only(nodes, &mut output);
4620            }
4621            ModuleOperationIr::MaxPool2dWithIndices(repr) => {
4622                repr.x.mark_read_only(nodes, &mut output);
4623            }
4624            ModuleOperationIr::MaxPool2dWithIndicesBackward(repr) => {
4625                repr.x.mark_read_only(nodes, &mut output);
4626                repr.grad.mark_read_only(nodes, &mut output);
4627            }
4628            ModuleOperationIr::AvgPool3d(repr) => {
4629                repr.x.mark_read_only(nodes, &mut output);
4630            }
4631            ModuleOperationIr::AvgPool3dBackward(repr) => {
4632                repr.x.mark_read_only(nodes, &mut output);
4633                repr.grad.mark_read_only(nodes, &mut output);
4634            }
4635            ModuleOperationIr::MaxPool3d(repr) => {
4636                repr.x.mark_read_only(nodes, &mut output);
4637            }
4638            ModuleOperationIr::MaxPool3dWithIndices(repr) => {
4639                repr.x.mark_read_only(nodes, &mut output);
4640            }
4641            ModuleOperationIr::MaxPool3dWithIndicesBackward(repr) => {
4642                repr.x.mark_read_only(nodes, &mut output);
4643                repr.grad.mark_read_only(nodes, &mut output);
4644            }
4645            ModuleOperationIr::Interpolate(repr) => {
4646                repr.x.mark_read_only(nodes, &mut output);
4647            }
4648            ModuleOperationIr::InterpolateBackward(repr) => {
4649                repr.x.mark_read_only(nodes, &mut output);
4650                repr.grad.mark_read_only(nodes, &mut output);
4651            }
4652            ModuleOperationIr::Attention(repr) => {
4653                repr.query.mark_read_only(nodes, &mut output);
4654                repr.key.mark_read_only(nodes, &mut output);
4655                repr.value.mark_read_only(nodes, &mut output);
4656                if let Some(mask) = &mut repr.mask {
4657                    mask.mark_read_only(nodes, &mut output);
4658                }
4659                if let Some(attn_bias) = &mut repr.attn_bias {
4660                    attn_bias.mark_read_only(nodes, &mut output);
4661                }
4662            }
4663            ModuleOperationIr::AttentionWithLogsumexp(repr) => {
4664                repr.query.mark_read_only(nodes, &mut output);
4665                repr.key.mark_read_only(nodes, &mut output);
4666                repr.value.mark_read_only(nodes, &mut output);
4667            }
4668            ModuleOperationIr::AttentionBackward(repr) => {
4669                repr.query.mark_read_only(nodes, &mut output);
4670                repr.key.mark_read_only(nodes, &mut output);
4671                repr.value.mark_read_only(nodes, &mut output);
4672                repr.out.mark_read_only(nodes, &mut output);
4673                repr.logsumexp.mark_read_only(nodes, &mut output);
4674                repr.out_grad.mark_read_only(nodes, &mut output);
4675            }
4676            ModuleOperationIr::CtcLoss(repr) => {
4677                repr.log_probs.mark_read_only(nodes, &mut output);
4678                repr.targets.mark_read_only(nodes, &mut output);
4679                repr.input_lengths.mark_read_only(nodes, &mut output);
4680                repr.target_lengths.mark_read_only(nodes, &mut output);
4681            }
4682            ModuleOperationIr::CtcLossBackward(repr) => {
4683                repr.log_probs.mark_read_only(nodes, &mut output);
4684                repr.targets.mark_read_only(nodes, &mut output);
4685                repr.input_lengths.mark_read_only(nodes, &mut output);
4686                repr.target_lengths.mark_read_only(nodes, &mut output);
4687                repr.grad_loss.mark_read_only(nodes, &mut output);
4688            }
4689            ModuleOperationIr::LayerNorm(repr) => {
4690                repr.input.mark_read_only(nodes, &mut output);
4691                repr.gamma.mark_read_only(nodes, &mut output);
4692                if let Some(beta) = &mut repr.beta {
4693                    beta.mark_read_only(nodes, &mut output);
4694                }
4695            }
4696            ModuleOperationIr::Unfold4d(repr) => {
4697                repr.x.mark_read_only(nodes, &mut output);
4698            }
4699            ModuleOperationIr::ConvTranspose1dWeightBackward(repr) => {
4700                repr.x.mark_read_only(nodes, &mut output);
4701                repr.weight.mark_read_only(nodes, &mut output);
4702                repr.output_grad.mark_read_only(nodes, &mut output);
4703            }
4704            ModuleOperationIr::ConvTranspose1dBiasBackward(repr) => {
4705                repr.x.mark_read_only(nodes, &mut output);
4706                repr.bias.mark_read_only(nodes, &mut output);
4707                repr.output_grad.mark_read_only(nodes, &mut output);
4708            }
4709            ModuleOperationIr::ConvTranspose2dWeightBackward(repr) => {
4710                repr.x.mark_read_only(nodes, &mut output);
4711                repr.weight.mark_read_only(nodes, &mut output);
4712                repr.output_grad.mark_read_only(nodes, &mut output);
4713            }
4714            ModuleOperationIr::ConvTranspose2dBiasBackward(repr) => {
4715                repr.x.mark_read_only(nodes, &mut output);
4716                repr.bias.mark_read_only(nodes, &mut output);
4717                repr.output_grad.mark_read_only(nodes, &mut output);
4718            }
4719            ModuleOperationIr::ConvTranspose3dWeightBackward(repr) => {
4720                repr.x.mark_read_only(nodes, &mut output);
4721                repr.weight.mark_read_only(nodes, &mut output);
4722                repr.output_grad.mark_read_only(nodes, &mut output);
4723            }
4724            ModuleOperationIr::ConvTranspose3dBiasBackward(repr) => {
4725                repr.x.mark_read_only(nodes, &mut output);
4726                repr.bias.mark_read_only(nodes, &mut output);
4727                repr.output_grad.mark_read_only(nodes, &mut output);
4728            }
4729        };
4730
4731        output
4732    }
4733
4734    fn visit_mut(&mut self, v: &mut impl IrVisitorMut) {
4735        match self {
4736            ModuleOperationIr::BatchNorm(repr) => {
4737                v.visit_tensor_mut(&mut repr.x);
4738                v.visit_tensor_mut(&mut repr.gamma);
4739                v.visit_tensor_mut(&mut repr.beta);
4740                v.visit_tensor_mut(&mut repr.mean);
4741                v.visit_tensor_mut(&mut repr.variance);
4742                v.visit_scalar_mut(&mut repr.epsilon);
4743                v.visit_tensor_mut(&mut repr.out);
4744            }
4745            ModuleOperationIr::Embedding(repr) => {
4746                v.visit_tensor_mut(&mut repr.weights);
4747                v.visit_tensor_mut(&mut repr.indices);
4748                v.visit_tensor_mut(&mut repr.out);
4749            }
4750            ModuleOperationIr::EmbeddingBackward(repr) => {
4751                v.visit_tensor_mut(&mut repr.weights);
4752                v.visit_tensor_mut(&mut repr.out_grad);
4753                v.visit_tensor_mut(&mut repr.indices);
4754                v.visit_tensor_mut(&mut repr.out);
4755            }
4756            ModuleOperationIr::Linear(repr) => {
4757                v.visit_tensor_mut(&mut repr.x);
4758                v.visit_tensor_mut(&mut repr.weight);
4759                if let Some(bias) = &mut repr.bias {
4760                    v.visit_tensor_mut(bias);
4761                }
4762                v.visit_tensor_mut(&mut repr.out);
4763            }
4764            ModuleOperationIr::LinearXBackward(repr) => {
4765                v.visit_tensor_mut(&mut repr.weight);
4766                v.visit_tensor_mut(&mut repr.output_grad);
4767                v.visit_tensor_mut(&mut repr.out);
4768            }
4769            ModuleOperationIr::LinearWeightBackward(repr) => {
4770                v.visit_tensor_mut(&mut repr.x);
4771                v.visit_tensor_mut(&mut repr.output_grad);
4772                v.visit_tensor_mut(&mut repr.out);
4773            }
4774            ModuleOperationIr::LinearBiasBackward(repr) => {
4775                v.visit_tensor_mut(&mut repr.output_grad);
4776                v.visit_tensor_mut(&mut repr.out);
4777            }
4778            ModuleOperationIr::Conv1d(repr) => {
4779                v.visit_tensor_mut(&mut repr.x);
4780                v.visit_tensor_mut(&mut repr.weight);
4781                if let Some(bias) = &mut repr.bias {
4782                    v.visit_tensor_mut(bias);
4783                }
4784                v.visit_tensor_mut(&mut repr.out);
4785            }
4786            ModuleOperationIr::Conv1dXBackward(repr) => {
4787                v.visit_tensor_mut(&mut repr.x);
4788                v.visit_tensor_mut(&mut repr.weight);
4789                v.visit_tensor_mut(&mut repr.output_grad);
4790                v.visit_tensor_mut(&mut repr.out);
4791            }
4792            ModuleOperationIr::Conv1dWeightBackward(repr) => {
4793                v.visit_tensor_mut(&mut repr.x);
4794                v.visit_tensor_mut(&mut repr.weight);
4795                v.visit_tensor_mut(&mut repr.output_grad);
4796                v.visit_tensor_mut(&mut repr.out);
4797            }
4798            ModuleOperationIr::Conv1dBiasBackward(repr) => {
4799                v.visit_tensor_mut(&mut repr.x);
4800                v.visit_tensor_mut(&mut repr.bias);
4801                v.visit_tensor_mut(&mut repr.output_grad);
4802                v.visit_tensor_mut(&mut repr.out);
4803            }
4804            ModuleOperationIr::Conv2d(repr) => {
4805                v.visit_tensor_mut(&mut repr.x);
4806                v.visit_tensor_mut(&mut repr.weight);
4807                if let Some(bias) = &mut repr.bias {
4808                    v.visit_tensor_mut(bias);
4809                }
4810                v.visit_tensor_mut(&mut repr.out);
4811            }
4812            ModuleOperationIr::Conv2dXBackward(repr) => {
4813                v.visit_tensor_mut(&mut repr.x);
4814                v.visit_tensor_mut(&mut repr.weight);
4815                v.visit_tensor_mut(&mut repr.output_grad);
4816                v.visit_tensor_mut(&mut repr.out);
4817            }
4818            ModuleOperationIr::Conv2dWeightBackward(repr) => {
4819                v.visit_tensor_mut(&mut repr.x);
4820                v.visit_tensor_mut(&mut repr.weight);
4821                v.visit_tensor_mut(&mut repr.output_grad);
4822                v.visit_tensor_mut(&mut repr.out);
4823            }
4824            ModuleOperationIr::Conv2dBiasBackward(repr) => {
4825                v.visit_tensor_mut(&mut repr.x);
4826                v.visit_tensor_mut(&mut repr.bias);
4827                v.visit_tensor_mut(&mut repr.output_grad);
4828                v.visit_tensor_mut(&mut repr.out);
4829            }
4830            ModuleOperationIr::Conv3d(repr) => {
4831                v.visit_tensor_mut(&mut repr.x);
4832                v.visit_tensor_mut(&mut repr.weight);
4833                if let Some(bias) = &mut repr.bias {
4834                    v.visit_tensor_mut(bias);
4835                }
4836                v.visit_tensor_mut(&mut repr.out);
4837            }
4838            ModuleOperationIr::Conv3dXBackward(repr) => {
4839                v.visit_tensor_mut(&mut repr.x);
4840                v.visit_tensor_mut(&mut repr.weight);
4841                v.visit_tensor_mut(&mut repr.output_grad);
4842                v.visit_tensor_mut(&mut repr.out);
4843            }
4844            ModuleOperationIr::Conv3dWeightBackward(repr) => {
4845                v.visit_tensor_mut(&mut repr.x);
4846                v.visit_tensor_mut(&mut repr.weight);
4847                v.visit_tensor_mut(&mut repr.output_grad);
4848                v.visit_tensor_mut(&mut repr.out);
4849            }
4850            ModuleOperationIr::Conv3dBiasBackward(repr) => {
4851                v.visit_tensor_mut(&mut repr.x);
4852                v.visit_tensor_mut(&mut repr.bias);
4853                v.visit_tensor_mut(&mut repr.output_grad);
4854                v.visit_tensor_mut(&mut repr.out);
4855            }
4856            ModuleOperationIr::DeformableConv2d(repr) => {
4857                v.visit_tensor_mut(&mut repr.x);
4858                v.visit_tensor_mut(&mut repr.offset);
4859                v.visit_tensor_mut(&mut repr.weight);
4860                if let Some(mask) = &mut repr.mask {
4861                    v.visit_tensor_mut(mask);
4862                }
4863                if let Some(bias) = &mut repr.bias {
4864                    v.visit_tensor_mut(bias);
4865                }
4866                v.visit_tensor_mut(&mut repr.out);
4867            }
4868            ModuleOperationIr::DeformableConv2dBackward(repr) => {
4869                v.visit_tensor_mut(&mut repr.x);
4870                v.visit_tensor_mut(&mut repr.offset);
4871                v.visit_tensor_mut(&mut repr.weight);
4872                v.visit_tensor_mut(&mut repr.out_grad);
4873                if let Some(mask) = &mut repr.mask {
4874                    v.visit_tensor_mut(mask);
4875                }
4876                if let Some(bias) = &mut repr.bias {
4877                    v.visit_tensor_mut(bias);
4878                }
4879                v.visit_tensor_mut(&mut repr.input_grad);
4880                v.visit_tensor_mut(&mut repr.offset_grad);
4881                v.visit_tensor_mut(&mut repr.weight_grad);
4882                if let Some(mask_grad) = &mut repr.mask_grad {
4883                    v.visit_tensor_mut(mask_grad);
4884                }
4885                if let Some(bias_grad) = &mut repr.bias_grad {
4886                    v.visit_tensor_mut(bias_grad);
4887                }
4888            }
4889            ModuleOperationIr::ConvTranspose1d(repr) => {
4890                v.visit_tensor_mut(&mut repr.x);
4891                v.visit_tensor_mut(&mut repr.weight);
4892                if let Some(bias) = &mut repr.bias {
4893                    v.visit_tensor_mut(bias);
4894                }
4895                v.visit_tensor_mut(&mut repr.out);
4896            }
4897            ModuleOperationIr::ConvTranspose2d(repr) => {
4898                v.visit_tensor_mut(&mut repr.x);
4899                v.visit_tensor_mut(&mut repr.weight);
4900                if let Some(bias) = &mut repr.bias {
4901                    v.visit_tensor_mut(bias);
4902                }
4903                v.visit_tensor_mut(&mut repr.out);
4904            }
4905            ModuleOperationIr::ConvTranspose3d(repr) => {
4906                v.visit_tensor_mut(&mut repr.x);
4907                v.visit_tensor_mut(&mut repr.weight);
4908                if let Some(bias) = &mut repr.bias {
4909                    v.visit_tensor_mut(bias);
4910                }
4911                v.visit_tensor_mut(&mut repr.out);
4912            }
4913            ModuleOperationIr::AvgPool1d(repr) => {
4914                v.visit_tensor_mut(&mut repr.x);
4915                v.visit_tensor_mut(&mut repr.out);
4916            }
4917            ModuleOperationIr::AvgPool2d(repr) => {
4918                v.visit_tensor_mut(&mut repr.x);
4919                v.visit_tensor_mut(&mut repr.out);
4920            }
4921            ModuleOperationIr::AvgPool1dBackward(repr) => {
4922                v.visit_tensor_mut(&mut repr.x);
4923                v.visit_tensor_mut(&mut repr.grad);
4924                v.visit_tensor_mut(&mut repr.out);
4925            }
4926            ModuleOperationIr::AvgPool2dBackward(repr) => {
4927                v.visit_tensor_mut(&mut repr.x);
4928                v.visit_tensor_mut(&mut repr.grad);
4929                v.visit_tensor_mut(&mut repr.out);
4930            }
4931            ModuleOperationIr::AdaptiveAvgPool1d(repr) => {
4932                v.visit_tensor_mut(&mut repr.x);
4933                v.visit_tensor_mut(&mut repr.out);
4934            }
4935            ModuleOperationIr::AdaptiveAvgPool2d(repr) => {
4936                v.visit_tensor_mut(&mut repr.x);
4937                v.visit_tensor_mut(&mut repr.out);
4938            }
4939            ModuleOperationIr::AdaptiveAvgPool1dBackward(repr) => {
4940                v.visit_tensor_mut(&mut repr.x);
4941                v.visit_tensor_mut(&mut repr.grad);
4942                v.visit_tensor_mut(&mut repr.out);
4943            }
4944            ModuleOperationIr::AdaptiveAvgPool2dBackward(repr) => {
4945                v.visit_tensor_mut(&mut repr.x);
4946                v.visit_tensor_mut(&mut repr.grad);
4947                v.visit_tensor_mut(&mut repr.out);
4948            }
4949            ModuleOperationIr::AdaptiveAvgPool3d(repr) => {
4950                v.visit_tensor_mut(&mut repr.x);
4951                v.visit_tensor_mut(&mut repr.out);
4952            }
4953            ModuleOperationIr::AdaptiveAvgPool3dBackward(repr) => {
4954                v.visit_tensor_mut(&mut repr.x);
4955                v.visit_tensor_mut(&mut repr.grad);
4956                v.visit_tensor_mut(&mut repr.out);
4957            }
4958            ModuleOperationIr::MaxPool1d(repr) => {
4959                v.visit_tensor_mut(&mut repr.x);
4960                v.visit_tensor_mut(&mut repr.out);
4961            }
4962            ModuleOperationIr::MaxPool1dWithIndices(repr) => {
4963                v.visit_tensor_mut(&mut repr.x);
4964                v.visit_tensor_mut(&mut repr.out);
4965                v.visit_tensor_mut(&mut repr.out_indices);
4966            }
4967            ModuleOperationIr::MaxPool1dWithIndicesBackward(repr) => {
4968                v.visit_tensor_mut(&mut repr.x);
4969                v.visit_tensor_mut(&mut repr.indices);
4970                v.visit_tensor_mut(&mut repr.grad);
4971                v.visit_tensor_mut(&mut repr.out);
4972            }
4973            ModuleOperationIr::MaxPool2d(repr) => {
4974                v.visit_tensor_mut(&mut repr.x);
4975                v.visit_tensor_mut(&mut repr.out);
4976            }
4977            ModuleOperationIr::MaxPool2dWithIndices(repr) => {
4978                v.visit_tensor_mut(&mut repr.x);
4979                v.visit_tensor_mut(&mut repr.out);
4980                v.visit_tensor_mut(&mut repr.out_indices);
4981            }
4982            ModuleOperationIr::MaxPool2dWithIndicesBackward(repr) => {
4983                v.visit_tensor_mut(&mut repr.x);
4984                v.visit_tensor_mut(&mut repr.indices);
4985                v.visit_tensor_mut(&mut repr.grad);
4986                v.visit_tensor_mut(&mut repr.out);
4987            }
4988            ModuleOperationIr::AvgPool3d(repr) => {
4989                v.visit_tensor_mut(&mut repr.x);
4990                v.visit_tensor_mut(&mut repr.out);
4991            }
4992            ModuleOperationIr::AvgPool3dBackward(repr) => {
4993                v.visit_tensor_mut(&mut repr.x);
4994                v.visit_tensor_mut(&mut repr.grad);
4995                v.visit_tensor_mut(&mut repr.out);
4996            }
4997            ModuleOperationIr::MaxPool3d(repr) => {
4998                v.visit_tensor_mut(&mut repr.x);
4999                v.visit_tensor_mut(&mut repr.out);
5000            }
5001            ModuleOperationIr::MaxPool3dWithIndices(repr) => {
5002                v.visit_tensor_mut(&mut repr.x);
5003                v.visit_tensor_mut(&mut repr.out);
5004                v.visit_tensor_mut(&mut repr.out_indices);
5005            }
5006            ModuleOperationIr::MaxPool3dWithIndicesBackward(repr) => {
5007                v.visit_tensor_mut(&mut repr.x);
5008                v.visit_tensor_mut(&mut repr.indices);
5009                v.visit_tensor_mut(&mut repr.grad);
5010                v.visit_tensor_mut(&mut repr.out);
5011            }
5012            ModuleOperationIr::Interpolate(repr) => {
5013                v.visit_tensor_mut(&mut repr.x);
5014                v.visit_tensor_mut(&mut repr.out);
5015            }
5016            ModuleOperationIr::InterpolateBackward(repr) => {
5017                v.visit_tensor_mut(&mut repr.x);
5018                v.visit_tensor_mut(&mut repr.grad);
5019                v.visit_tensor_mut(&mut repr.out);
5020            }
5021            ModuleOperationIr::Attention(repr) => {
5022                v.visit_tensor_mut(&mut repr.query);
5023                v.visit_tensor_mut(&mut repr.key);
5024                v.visit_tensor_mut(&mut repr.value);
5025                if let Some(mask) = &mut repr.mask {
5026                    v.visit_tensor_mut(mask);
5027                }
5028                if let Some(attn_bias) = &mut repr.attn_bias {
5029                    v.visit_tensor_mut(attn_bias);
5030                }
5031                v.visit_tensor_mut(&mut repr.out);
5032                if let Some(scale) = &mut repr.options.scale {
5033                    v.visit_scalar_mut(scale);
5034                }
5035                if let Some(softcap) = &mut repr.options.softcap {
5036                    v.visit_scalar_mut(softcap);
5037                }
5038            }
5039            ModuleOperationIr::AttentionWithLogsumexp(repr) => {
5040                v.visit_tensor_mut(&mut repr.query);
5041                v.visit_tensor_mut(&mut repr.key);
5042                v.visit_tensor_mut(&mut repr.value);
5043                v.visit_tensor_mut(&mut repr.out);
5044                v.visit_tensor_mut(&mut repr.logsumexp);
5045                if let Some(scale) = &mut repr.options.scale {
5046                    v.visit_scalar_mut(scale);
5047                }
5048                if let Some(softcap) = &mut repr.options.softcap {
5049                    v.visit_scalar_mut(softcap);
5050                }
5051            }
5052            ModuleOperationIr::AttentionBackward(repr) => {
5053                v.visit_tensor_mut(&mut repr.query);
5054                v.visit_tensor_mut(&mut repr.key);
5055                v.visit_tensor_mut(&mut repr.value);
5056                v.visit_tensor_mut(&mut repr.out);
5057                v.visit_tensor_mut(&mut repr.logsumexp);
5058                v.visit_tensor_mut(&mut repr.out_grad);
5059                v.visit_tensor_mut(&mut repr.query_grad);
5060                v.visit_tensor_mut(&mut repr.key_grad);
5061                v.visit_tensor_mut(&mut repr.value_grad);
5062                if let Some(scale) = &mut repr.options.scale {
5063                    v.visit_scalar_mut(scale);
5064                }
5065                if let Some(softcap) = &mut repr.options.softcap {
5066                    v.visit_scalar_mut(softcap);
5067                }
5068            }
5069            ModuleOperationIr::CtcLoss(repr) => {
5070                v.visit_tensor_mut(&mut repr.log_probs);
5071                v.visit_tensor_mut(&mut repr.targets);
5072                v.visit_tensor_mut(&mut repr.input_lengths);
5073                v.visit_tensor_mut(&mut repr.target_lengths);
5074                v.visit_tensor_mut(&mut repr.out);
5075            }
5076            ModuleOperationIr::CtcLossBackward(repr) => {
5077                v.visit_tensor_mut(&mut repr.log_probs);
5078                v.visit_tensor_mut(&mut repr.targets);
5079                v.visit_tensor_mut(&mut repr.input_lengths);
5080                v.visit_tensor_mut(&mut repr.target_lengths);
5081                v.visit_tensor_mut(&mut repr.grad_loss);
5082                v.visit_tensor_mut(&mut repr.out);
5083            }
5084            ModuleOperationIr::LayerNorm(repr) => {
5085                v.visit_tensor_mut(&mut repr.input);
5086                v.visit_tensor_mut(&mut repr.gamma);
5087                if let Some(beta) = &mut repr.beta {
5088                    v.visit_tensor_mut(beta);
5089                }
5090                v.visit_tensor_mut(&mut repr.out);
5091                v.visit_scalar_mut(&mut repr.epsilon);
5092            }
5093            ModuleOperationIr::Unfold4d(repr) => {
5094                v.visit_tensor_mut(&mut repr.x);
5095                v.visit_tensor_mut(&mut repr.out);
5096            }
5097            ModuleOperationIr::ConvTranspose1dWeightBackward(repr) => {
5098                v.visit_tensor_mut(&mut repr.x);
5099                v.visit_tensor_mut(&mut repr.weight);
5100                v.visit_tensor_mut(&mut repr.output_grad);
5101                v.visit_tensor_mut(&mut repr.out);
5102            }
5103            ModuleOperationIr::ConvTranspose1dBiasBackward(repr) => {
5104                v.visit_tensor_mut(&mut repr.x);
5105                v.visit_tensor_mut(&mut repr.bias);
5106                v.visit_tensor_mut(&mut repr.output_grad);
5107                v.visit_tensor_mut(&mut repr.out);
5108            }
5109            ModuleOperationIr::ConvTranspose2dWeightBackward(repr) => {
5110                v.visit_tensor_mut(&mut repr.x);
5111                v.visit_tensor_mut(&mut repr.weight);
5112                v.visit_tensor_mut(&mut repr.output_grad);
5113                v.visit_tensor_mut(&mut repr.out);
5114            }
5115            ModuleOperationIr::ConvTranspose2dBiasBackward(repr) => {
5116                v.visit_tensor_mut(&mut repr.x);
5117                v.visit_tensor_mut(&mut repr.bias);
5118                v.visit_tensor_mut(&mut repr.output_grad);
5119                v.visit_tensor_mut(&mut repr.out);
5120            }
5121            ModuleOperationIr::ConvTranspose3dWeightBackward(repr) => {
5122                v.visit_tensor_mut(&mut repr.x);
5123                v.visit_tensor_mut(&mut repr.weight);
5124                v.visit_tensor_mut(&mut repr.output_grad);
5125                v.visit_tensor_mut(&mut repr.out);
5126            }
5127            ModuleOperationIr::ConvTranspose3dBiasBackward(repr) => {
5128                v.visit_tensor_mut(&mut repr.x);
5129                v.visit_tensor_mut(&mut repr.bias);
5130                v.visit_tensor_mut(&mut repr.output_grad);
5131                v.visit_tensor_mut(&mut repr.out);
5132            }
5133        }
5134    }
5135}
5136
5137impl InitOperationIr {
5138    fn inputs(&self) -> Box<dyn Iterator<Item = &TensorIr> + '_> {
5139        Box::new([].into_iter())
5140    }
5141    fn outputs(&self) -> Box<dyn Iterator<Item = &TensorIr> + '_> {
5142        Box::new([&self.out].into_iter())
5143    }
5144
5145    fn visit_mut(&mut self, v: &mut impl IrVisitorMut) {
5146        v.visit_tensor_mut(&mut self.out);
5147    }
5148}
5149
5150impl TensorIr {
5151    fn mark_read_only(&mut self, nodes: &[TensorId], output: &mut Vec<TensorIr>) {
5152        if self.status == TensorStatus::ReadWrite && nodes.contains(&self.id) {
5153            output.push(self.clone());
5154            self.status = TensorStatus::ReadOnly;
5155        }
5156    }
5157}
5158
5159impl core::hash::Hash for RandomOpIr {
5160    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
5161        self.out.hash(state);
5162
5163        match self.distribution {
5164            Distribution::Default => 1u8.hash(state),
5165            Distribution::Bernoulli(_) => 2u8.hash(state),
5166            Distribution::Uniform(_, _) => 3u8.hash(state),
5167            Distribution::Normal(_, _) => 4u8.hash(state),
5168        }
5169    }
5170}
5171
5172/// Extension trait to extract outputs when registering an operation.
5173pub trait OperationOutput<O> {
5174    /// Extract a single output.
5175    fn output(self) -> O;
5176
5177    /// Extract a fixed number of outputs.
5178    fn outputs<const N: usize>(self) -> [O; N];
5179}
5180
5181impl<O: core::fmt::Debug> OperationOutput<O> for Vec<O> {
5182    fn output(self) -> O {
5183        let [tensor] = self.outputs();
5184        tensor
5185    }
5186
5187    fn outputs<const N: usize>(self) -> [O; N] {
5188        self.try_into().unwrap()
5189    }
5190}
5191
5192/// Operation IR for sort along a dim. The output preserves the input shape/dtype.
5193#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
5194pub struct SortOpIr {
5195    /// Input tensor.
5196    pub input: TensorIr,
5197    /// Dim along which to sort.
5198    pub dim: usize,
5199    /// Sort descending.
5200    pub descending: bool,
5201    /// Output tensor (same shape/dtype as input).
5202    pub out: TensorIr,
5203}
5204
5205/// Operation IR for sort-with-indices: returns sorted values + source indices.
5206#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
5207pub struct SortWithIndicesOpIr {
5208    /// Input tensor.
5209    pub input: TensorIr,
5210    /// Dim along which to sort.
5211    pub dim: usize,
5212    /// Sort descending.
5213    pub descending: bool,
5214    /// Output tensor with sorted values.
5215    pub out: TensorIr,
5216    /// Output tensor with the indices into the original input.
5217    pub out_indices: TensorIr,
5218}
5219
5220/// Operation IR for layer normalization with optional bias.
5221#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
5222pub struct LayerNormOpIr {
5223    /// Input tensor.
5224    pub input: TensorIr,
5225    /// Scale (gamma) parameter.
5226    pub gamma: TensorIr,
5227    /// Optional shift (beta) parameter.
5228    pub beta: Option<TensorIr>,
5229    /// Numerical-stability epsilon.
5230    pub epsilon: ScalarIr,
5231    /// Output tensor.
5232    pub out: TensorIr,
5233}
5234
5235/// Operation IR for unfold4d (a 4d sliding-window kernel-as-explicit-columns reshape).
5236#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
5237pub struct Unfold4dOpIr {
5238    /// Input tensor of shape `[N, C, H, W]`.
5239    pub x: TensorIr,
5240    /// Kernel size `[kH, kW]`.
5241    pub kernel_size: [usize; 2],
5242    /// Conv-like options (stride, padding, dilation).
5243    pub options: Unfold4dOptionsIr,
5244    /// Output tensor.
5245    pub out: TensorIr,
5246}
5247
5248/// Options for [`Unfold4dOpIr`] — mirrors the backend's `UnfoldOptions`.
5249#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
5250pub struct Unfold4dOptionsIr {
5251    /// Stride `[sH, sW]`.
5252    pub stride: [usize; 2],
5253    /// Padding `[pH, pW]`.
5254    pub padding: [usize; 2],
5255    /// Dilation `[dH, dW]`.
5256    pub dilation: [usize; 2],
5257}
5258
5259/// Operation IR for the weight-gradient of `conv_transpose1d`.
5260#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
5261pub struct ConvTranspose1dWeightBackwardOpIr {
5262    /// Forward input.
5263    pub x: TensorIr,
5264    /// Convolution weights.
5265    pub weight: TensorIr,
5266    /// Upstream gradient.
5267    pub output_grad: TensorIr,
5268    /// Convolution options.
5269    pub options: ConvTranspose1dOptionsIr,
5270    /// Output: gradient w.r.t. `weight`.
5271    pub out: TensorIr,
5272}
5273
5274/// Operation IR for the bias-gradient of `conv_transpose1d`.
5275#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
5276pub struct ConvTranspose1dBiasBackwardOpIr {
5277    /// Forward input.
5278    pub x: TensorIr,
5279    /// Bias tensor (sets the output shape).
5280    pub bias: TensorIr,
5281    /// Upstream gradient.
5282    pub output_grad: TensorIr,
5283    /// Output: gradient w.r.t. `bias`.
5284    pub out: TensorIr,
5285}
5286
5287/// Operation IR for the weight-gradient of `conv_transpose2d`.
5288#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
5289pub struct ConvTranspose2dWeightBackwardOpIr {
5290    /// Forward input.
5291    pub x: TensorIr,
5292    /// Convolution weights.
5293    pub weight: TensorIr,
5294    /// Upstream gradient.
5295    pub output_grad: TensorIr,
5296    /// Convolution options.
5297    pub options: ConvTranspose2dOptionsIr,
5298    /// Output: gradient w.r.t. `weight`.
5299    pub out: TensorIr,
5300}
5301
5302/// Operation IR for the bias-gradient of `conv_transpose2d`.
5303#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
5304pub struct ConvTranspose2dBiasBackwardOpIr {
5305    /// Forward input.
5306    pub x: TensorIr,
5307    /// Bias tensor.
5308    pub bias: TensorIr,
5309    /// Upstream gradient.
5310    pub output_grad: TensorIr,
5311    /// Output: gradient w.r.t. `bias`.
5312    pub out: TensorIr,
5313}
5314
5315/// Operation IR for the weight-gradient of `conv_transpose3d`.
5316#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
5317pub struct ConvTranspose3dWeightBackwardOpIr {
5318    /// Forward input.
5319    pub x: TensorIr,
5320    /// Convolution weights.
5321    pub weight: TensorIr,
5322    /// Upstream gradient.
5323    pub output_grad: TensorIr,
5324    /// Convolution options.
5325    pub options: ConvTranspose3dOptionsIr,
5326    /// Output: gradient w.r.t. `weight`.
5327    pub out: TensorIr,
5328}
5329
5330/// Operation IR for the bias-gradient of `conv_transpose3d`.
5331#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
5332pub struct ConvTranspose3dBiasBackwardOpIr {
5333    /// Forward input.
5334    pub x: TensorIr,
5335    /// Bias tensor.
5336    pub bias: TensorIr,
5337    /// Upstream gradient.
5338    pub output_grad: TensorIr,
5339    /// Output: gradient w.r.t. `bias`.
5340    pub out: TensorIr,
5341}
5342
5343/// Operation IR for the hard-sigmoid activation function, which takes two scalars
5344/// (`alpha` and `beta`) in addition to the input tensor.
5345#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
5346pub struct HardSigmoidOpIr {
5347    /// Input tensor.
5348    pub tensor: TensorIr,
5349    /// Alpha — multiplied with the tensor before adding `beta`.
5350    pub alpha: ScalarIr,
5351    /// Beta — added after multiplying by `alpha`.
5352    pub beta: ScalarIr,
5353    /// Output tensor.
5354    pub out: TensorIr,
5355}
5356
5357/// Operation intermediate representation for activation functions.
5358///
5359/// Activations are kept in their own enum (rather than smeared across `FloatOperationIr`)
5360/// so backends can match a single arm to dispatch the fused implementation, instead of
5361/// receiving the function decomposed into 5–20 primitive ops.
5362#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)]
5363pub enum ActivationOperationIr {
5364    /// [relu](burn_backend::ops::ActivationOps::relu).
5365    Relu(UnaryOpIr),
5366    /// [relu_backward](burn_backend::ops::ActivationOps::relu_backward).
5367    /// `lhs` = output of the forward pass, `rhs` = upstream gradient.
5368    ReluBackward(BinaryOpIr),
5369    /// [leaky_relu](burn_backend::ops::ActivationOps::leaky_relu).
5370    /// `lhs` = input, `rhs` = `negative_slope` scalar.
5371    LeakyRelu(ScalarOpIr),
5372    /// [prelu](burn_backend::ops::ActivationOps::prelu).
5373    /// `lhs` = input, `rhs` = alpha tensor.
5374    PRelu(BinaryOpIr),
5375    /// [gelu](burn_backend::ops::ActivationOps::gelu).
5376    Gelu(UnaryOpIr),
5377    /// [gelu_backward](burn_backend::ops::ActivationOps::gelu_backward).
5378    /// `lhs` = forward input, `rhs` = upstream gradient.
5379    GeluBackward(BinaryOpIr),
5380    /// [sigmoid](burn_backend::ops::ActivationOps::sigmoid).
5381    Sigmoid(UnaryOpIr),
5382    /// [sigmoid_backward](burn_backend::ops::ActivationOps::sigmoid_backward).
5383    /// `lhs` = output of the forward pass, `rhs` = upstream gradient.
5384    SigmoidBackward(BinaryOpIr),
5385    /// [hard_sigmoid](burn_backend::ops::ActivationOps::hard_sigmoid).
5386    HardSigmoid(HardSigmoidOpIr),
5387    /// [log_sigmoid](burn_backend::ops::ActivationOps::log_sigmoid).
5388    LogSigmoid(UnaryOpIr),
5389    /// [log_sigmoid_backward](burn_backend::ops::ActivationOps::log_sigmoid_backward).
5390    /// `lhs` = forward input, `rhs` = upstream gradient.
5391    LogSigmoidBackward(BinaryOpIr),
5392    /// [softmax](burn_backend::ops::ActivationOps::softmax).
5393    Softmax(DimOpIr),
5394    /// [log_softmax](burn_backend::ops::ActivationOps::log_softmax).
5395    LogSoftmax(DimOpIr),
5396    /// [softmin](burn_backend::ops::ActivationOps::softmin).
5397    Softmin(DimOpIr),
5398}
5399
5400/// Generate [`ActivationOperationIr`]'s `inputs`/`outputs`/`mark_read_only` from a single
5401/// variant → input-field table, so the three methods can't drift out of sync. The `out`
5402/// field is implicit (every variant's op IR has one). Each entry lists the variant and the
5403/// names of its input tensor fields (`input`, `lhs`/`rhs`, `tensor`, …).
5404macro_rules! activation_ir_tensor_access {
5405    ($( $variant:ident => [ $($field:ident),+ ] ),+ $(,)?) => {
5406        impl ActivationOperationIr {
5407            fn inputs(&self) -> Box<dyn Iterator<Item = &TensorIr> + '_> {
5408                match self {
5409                    $( Self::$variant(repr) => Box::new([$(&repr.$field),+].into_iter()), )+
5410                }
5411            }
5412
5413            fn outputs(&self) -> Box<dyn Iterator<Item = &TensorIr> + '_> {
5414                match self {
5415                    $( Self::$variant(repr) => Box::new([&repr.out].into_iter()), )+
5416                }
5417            }
5418
5419            fn mark_read_only(&mut self, nodes: &[TensorId]) -> Vec<TensorIr> {
5420                let mut output = Vec::new();
5421                match self {
5422                    $( Self::$variant(repr) => {
5423                        $( repr.$field.mark_read_only(nodes, &mut output); )+
5424                    } )+
5425                }
5426                output
5427            }
5428        }
5429    };
5430}
5431
5432impl ActivationOperationIr {
5433    fn visit_mut(&mut self, v: &mut impl IrVisitorMut) {
5434        match self {
5435            ActivationOperationIr::Relu(repr) => {
5436                v.visit_tensor_mut(&mut repr.input);
5437                v.visit_tensor_mut(&mut repr.out);
5438            }
5439            ActivationOperationIr::ReluBackward(repr) => {
5440                v.visit_tensor_mut(&mut repr.lhs);
5441                v.visit_tensor_mut(&mut repr.rhs);
5442                v.visit_tensor_mut(&mut repr.out);
5443            }
5444            ActivationOperationIr::LeakyRelu(repr) => {
5445                v.visit_tensor_mut(&mut repr.lhs);
5446                v.visit_tensor_mut(&mut repr.out);
5447                v.visit_scalar_mut(&mut repr.rhs);
5448            }
5449            ActivationOperationIr::PRelu(repr) => {
5450                v.visit_tensor_mut(&mut repr.lhs);
5451                v.visit_tensor_mut(&mut repr.rhs);
5452                v.visit_tensor_mut(&mut repr.out);
5453            }
5454            ActivationOperationIr::Gelu(repr) => {
5455                v.visit_tensor_mut(&mut repr.input);
5456                v.visit_tensor_mut(&mut repr.out);
5457            }
5458            ActivationOperationIr::GeluBackward(repr) => {
5459                v.visit_tensor_mut(&mut repr.lhs);
5460                v.visit_tensor_mut(&mut repr.rhs);
5461                v.visit_tensor_mut(&mut repr.out);
5462            }
5463            ActivationOperationIr::Sigmoid(repr) => {
5464                v.visit_tensor_mut(&mut repr.input);
5465                v.visit_tensor_mut(&mut repr.out);
5466            }
5467            ActivationOperationIr::SigmoidBackward(repr) => {
5468                v.visit_tensor_mut(&mut repr.lhs);
5469                v.visit_tensor_mut(&mut repr.rhs);
5470                v.visit_tensor_mut(&mut repr.out);
5471            }
5472            ActivationOperationIr::HardSigmoid(repr) => {
5473                v.visit_tensor_mut(&mut repr.tensor);
5474                v.visit_tensor_mut(&mut repr.out);
5475                v.visit_scalar_mut(&mut repr.alpha);
5476                v.visit_scalar_mut(&mut repr.beta);
5477            }
5478            ActivationOperationIr::LogSigmoid(repr) => {
5479                v.visit_tensor_mut(&mut repr.input);
5480                v.visit_tensor_mut(&mut repr.out);
5481            }
5482            ActivationOperationIr::LogSigmoidBackward(repr) => {
5483                v.visit_tensor_mut(&mut repr.lhs);
5484                v.visit_tensor_mut(&mut repr.rhs);
5485                v.visit_tensor_mut(&mut repr.out);
5486            }
5487            ActivationOperationIr::Softmax(repr) => {
5488                v.visit_tensor_mut(&mut repr.input);
5489                v.visit_tensor_mut(&mut repr.out);
5490            }
5491            ActivationOperationIr::LogSoftmax(repr) => {
5492                v.visit_tensor_mut(&mut repr.input);
5493                v.visit_tensor_mut(&mut repr.out);
5494            }
5495            ActivationOperationIr::Softmin(repr) => {
5496                v.visit_tensor_mut(&mut repr.input);
5497                v.visit_tensor_mut(&mut repr.out);
5498            }
5499        }
5500    }
5501}
5502
5503activation_ir_tensor_access! {
5504    Relu => [input],
5505    ReluBackward => [lhs, rhs],
5506    LeakyRelu => [lhs],
5507    PRelu => [lhs, rhs],
5508    Gelu => [input],
5509    GeluBackward => [lhs, rhs],
5510    Sigmoid => [input],
5511    SigmoidBackward => [lhs, rhs],
5512    HardSigmoid => [tensor],
5513    LogSigmoid => [input],
5514    LogSigmoidBackward => [lhs, rhs],
5515    Softmax => [input],
5516    LogSoftmax => [input],
5517    Softmin => [input],
5518}
5519
5520#[cfg(test)]
5521mod tests {
5522    use super::*;
5523    use burn_backend::{DType, Shape};
5524
5525    #[test]
5526    #[should_panic(expected = "expected symmetric convolution padding")]
5527    fn conv3d_options_ir_rejects_asymmetric_padding() {
5528        let options =
5529            ConvOptions::new_with_padding([1, 1, 1], [(0, 1), (2, 2), (3, 3)], [1, 1, 1], 1);
5530
5531        let _: Conv3dOptionsIr = options.into();
5532    }
5533
5534    fn tensor(id: u64) -> TensorIr {
5535        TensorIr::uninit(TensorId::new(id), Shape::from([2, 2]), DType::F32)
5536    }
5537
5538    /// Bumps every visited tensor id by 100 and collects (and rewrites) every visited scalar.
5539    #[derive(Default)]
5540    struct CollectVisitor {
5541        scalars: Vec<ScalarIr>,
5542        rewrite_scalar: Option<ScalarIr>,
5543    }
5544
5545    impl IrVisitorMut for CollectVisitor {
5546        fn visit_tensor_mut(&mut self, tensor: &mut TensorIr) {
5547            tensor.id = TensorId::new(tensor.id.value() + 100);
5548        }
5549
5550        fn visit_scalar_mut(&mut self, scalar: &mut ScalarIr) {
5551            self.scalars.push(*scalar);
5552            if let Some(value) = self.rewrite_scalar {
5553                *scalar = value;
5554            }
5555        }
5556    }
5557
5558    #[test]
5559    fn visit_mut_visits_all_tensors_and_scalars() {
5560        // NumericFloat MulScalar (a ScalarOpIr with a scalar `rhs`).
5561        let mut mul = OperationIr::NumericFloat(
5562            DType::F32,
5563            NumericOperationIr::MulScalar(ScalarOpIr {
5564                lhs: tensor(1),
5565                rhs: ScalarIr::Float(2.0),
5566                out: tensor(2),
5567            }),
5568        );
5569
5570        // Bump every tensor id by 100 and rewrite the scalar to 9.0.
5571        let mut visitor = CollectVisitor {
5572            rewrite_scalar: Some(ScalarIr::Float(9.0)),
5573            ..Default::default()
5574        };
5575        mul.visit_mut(&mut visitor);
5576
5577        let ids: Vec<u64> = mul
5578            .inputs()
5579            .chain(mul.outputs())
5580            .map(|t| t.id.value())
5581            .collect();
5582        assert_eq!(ids, vec![101, 102]);
5583        // The scalar must have been visited (and was rewritable).
5584        assert_eq!(visitor.scalars, vec![ScalarIr::Float(2.0)]);
5585
5586        // Visiting again observes the rewritten scalar.
5587        let mut after = CollectVisitor::default();
5588        mul.visit_mut(&mut after);
5589        assert_eq!(after.scalars, vec![ScalarIr::Float(9.0)]);
5590
5591        // BaseFloat Reshape (a ShapeOpIr, input + out, no scalars).
5592        let mut reshape = OperationIr::BaseFloat(BaseOperationIr::Reshape(ShapeOpIr {
5593            input: tensor(10),
5594            out: tensor(11),
5595        }));
5596
5597        let mut visitor = CollectVisitor::default();
5598        reshape.visit_mut(&mut visitor);
5599        let ids: Vec<u64> = reshape
5600            .inputs()
5601            .chain(reshape.outputs())
5602            .map(|t| t.id.value())
5603            .collect();
5604        assert_eq!(ids, vec![110, 111]);
5605        assert_eq!(visitor.scalars.len(), 0);
5606    }
5607
5608    #[test]
5609    fn pad_ir_tracks_shape_tensors_and_constant_value() {
5610        let input = TensorIr::uninit(TensorId::new(1), Shape::from([2, 3]), DType::F32);
5611        let desc = PadOpIr::create(
5612            input,
5613            vec![(1, 2), (3, 4)],
5614            PadModeIr::Constant(ScalarIr::Float(5.0)),
5615            || TensorId::new(2),
5616        );
5617        assert_eq!(desc.out.shape, Shape::from([5, 10]));
5618
5619        let mut op = OperationIr::NumericFloat(DType::F32, NumericOperationIr::Pad(desc));
5620        let mut visitor = CollectVisitor {
5621            rewrite_scalar: Some(ScalarIr::Float(7.0)),
5622            ..Default::default()
5623        };
5624        op.visit_mut(&mut visitor);
5625
5626        let ids = op
5627            .inputs()
5628            .chain(op.outputs())
5629            .map(|tensor| tensor.id.value())
5630            .collect::<Vec<_>>();
5631        assert_eq!(ids, vec![101, 102]);
5632        assert_eq!(visitor.scalars, vec![ScalarIr::Float(5.0)]);
5633
5634        let OperationIr::NumericFloat(_, NumericOperationIr::Pad(desc)) = op else {
5635            panic!("expected pad operation");
5636        };
5637        assert_eq!(desc.mode, PadModeIr::Constant(ScalarIr::Float(7.0)));
5638    }
5639
5640    #[test]
5641    fn interpolate_options_ir_hash_canonicalizes_signed_zero() {
5642        use core::hash::{Hash, Hasher};
5643        use std::collections::hash_map::DefaultHasher;
5644
5645        fn calculate_hash<T: Hash>(t: &T) -> u64 {
5646            let mut s = DefaultHasher::new();
5647            t.hash(&mut s);
5648            s.finish()
5649        }
5650
5651        let opt1 = InterpolateOptionsIr {
5652            mode: InterpolateModeIr::Bicubic,
5653            coordinate_transform: CoordinateTransformModeIr::TfCropAndResize {
5654                roi: [[0.0, 0.0], [1.0, 1.0]],
5655                extrapolation_value: 0.0,
5656            },
5657            nearest_mode: NearestRoundingModeIr::Floor,
5658            cubic_options: CubicOptionsIr {
5659                a: 0.0,
5660                exclude_outside: true,
5661            },
5662            output_size: Some([4, 4]),
5663            scale_factor: None,
5664        };
5665
5666        let opt2 = InterpolateOptionsIr {
5667            mode: InterpolateModeIr::Bicubic,
5668            coordinate_transform: CoordinateTransformModeIr::TfCropAndResize {
5669                roi: [[-0.0, -0.0], [1.0, 1.0]],
5670                extrapolation_value: -0.0,
5671            },
5672            nearest_mode: NearestRoundingModeIr::Floor,
5673            cubic_options: CubicOptionsIr {
5674                a: -0.0,
5675                exclude_outside: true,
5676            },
5677            output_size: Some([4, 4]),
5678            scale_factor: None,
5679        };
5680
5681        assert_eq!(opt1, opt2);
5682        assert_eq!(calculate_hash(&opt1), calculate_hash(&opt2));
5683    }
5684}
burn_ir/operation.rs

burn_ir/
operation.rs
