from __future__ import annotations

import re
from pathlib import Path
from typing import Any, Callable, Dict, Optional, cast

from ._pytorch_exporter_native_codegen_common import (
    _NativeCodegenBindings,
    _NativeCodegenState,
)


_NATIVE_CODEGEN_FUNCTION_SOURCE = 'def _write_native_model_file_codegen_core_body_main_inner_legacy_impl(\n    context: _NativeModelFileWriterContext,\n) -> List[Tuple[str, str]]:\n    output_folder_path = context.output_folder_path\n    model_ir = context.model_ir\n    metadata = context.metadata\n    tensor_storage_name_map = context.tensor_storage_name_map\n    package_dir = context.package_dir\n    _ensure_direct_codegen_supported(model_ir)\n    preserve_channel_last_tensor_names = context.preserve_channel_last_tensor_names\n    tensor_var_names = context.tensor_var_names\n    producer_index = context.producer_index\n    consumer_index = context.consumer_index\n    module_param_tensor_names: Set[str] = set()\n    submodule_state_tensor_names: Set[str] = set()\n    module_init_lines = context.module_init_lines\n    load_specs = context.load_specs\n    op_module_attr_names: Dict[int, str] = {}\n    fused_module_specs: Dict[int, Dict[str, Any]] = {}\n    affine_layer_norm_specs: Dict[int, Dict[str, Any]] = {}\n    nms_method_specs: List[Dict[str, Any]] = []\n    module_attr_counts: Dict[str, int] = {}\n    inlined_constant_tensor_names: Set[str] = set()\n    skipped_op_indices: Set[int] = set()\n    conv_module_pad_specs: Dict[int, Optional[List[int]]] = {}\n    tensor_expr_aliases: Dict[str, str] = {}\n    channel_first_tensor_expr_aliases: Dict[str, str] = {}\n    synthetic_local_var_names: Dict[str, str] = {}\n    used_local_var_names: Set[str] = set(tensor_var_names.values())\n    synthetic_tensor_serial_ref: List[int] = [0]\n    public_input_names: Set[str] = {str(name) for name in list(model_ir.inputs)}\n    public_layout_bridge_tensor_names: Set[str] = {\n        str(name)\n        for name in list(model_ir.metadata.get("public_layout_bridge_tensor_names", []))\n        if str(name) != ""\n    }\n\n    _fold_single_consumer_public_input_bridge = lambda *, tensor_name, downstream_permute: _fold_single_consumer_public_input_bridge_for_codegen(\n        model_ir=model_ir,\n        producer_index=producer_index,\n        consumer_index=consumer_index,\n        public_layout_bridge_tensor_names=public_layout_bridge_tensor_names,\n        public_input_names=public_input_names,\n        tensor_name=tensor_name,\n        downstream_permute=downstream_permute,\n    )\n    _match_single_consumer_layout_bridge_transpose = lambda *, tensor_name, required_output_layout=None: _match_single_consumer_layout_bridge_transpose_for_codegen(\n        model_ir=model_ir,\n        consumer_index=consumer_index,\n        tensor_name=tensor_name,\n        required_output_layout=required_output_layout,\n    )\n    _target_shape_values = lambda tensor_name: _target_shape_values_for_model_ir(\n        model_ir=model_ir,\n        tensor_name=tensor_name,\n    )\n    _target_shape_literal = lambda tensor_name: _target_shape_literal_for_model_ir(\n        model_ir=model_ir,\n        tensor_name=tensor_name,\n    )\n    _resize_target_shape_literal = lambda output_name, input_name: _resize_target_shape_literal_for_model_ir(\n        model_ir=model_ir,\n        output_name=output_name,\n        input_name=input_name,\n    )\n    _tensor_shape_list = lambda tensor_name: _tensor_shape_list_for_model_ir(\n        model_ir=model_ir,\n        tensor_name=tensor_name,\n    )\n    _rank4_channel_first_shape_for_tensor = lambda tensor_name: _rank4_channel_first_shape_for_tensor_for_codegen(\n        model_ir=model_ir,\n        channel_first_tensor_expr_aliases=channel_first_tensor_expr_aliases,\n        tensor_name=tensor_name,\n    )\n    _channel_first_shape_for_tensor = lambda tensor_name: _channel_first_shape_for_tensor_for_codegen(\n        model_ir=model_ir,\n        channel_first_tensor_expr_aliases=channel_first_tensor_expr_aliases,\n        tensor_name=tensor_name,\n    )\n\n    _channel_first_concat_input_expr = lambda tensor_name: _channel_first_concat_input_expr_for_codegen(\n        model_ir=model_ir,\n        channel_first_tensor_expr_aliases=channel_first_tensor_expr_aliases,\n        tensor_name=tensor_name,\n        tensor_expr_fn=_tensor_expr,\n    )\n    _can_fold_channel_last_alias_slice_consumer = lambda op, *, expected_input_name: _can_fold_channel_last_alias_slice_consumer_for_codegen(\n        model_ir=model_ir,\n        op=op,\n        expected_input_name=expected_input_name,\n    )\n    _resolve_concat_axis_for_channel_first = lambda op: _resolve_concat_axis_for_channel_first_for_codegen(\n        model_ir=model_ir,\n        op=op,\n        channel_first_shape_for_tensor_fn=_channel_first_shape_for_tensor,\n        tensor_shape_list_fn=_tensor_shape_list,\n    )\n    _can_keep_channel_first_slice_output = lambda output_name: _can_keep_channel_first_slice_output_for_codegen(\n        model_ir=model_ir,\n        consumer_index=consumer_index,\n        output_name=output_name,\n        resolve_concat_axis_for_channel_first_fn=_resolve_concat_axis_for_channel_first,\n    )\n\n    _reshape_codegen_is_plain_data_only = lambda op: _reshape_codegen_is_plain_data_only_for_codegen(\n        model_ir=model_ir,\n        op=op,\n        infer_effective_rank4_runtime_layout_fn=_infer_effective_rank4_runtime_layout,\n        reshape_preserves_channel_last_sequence_fn=_reshape_preserves_channel_last_sequence,\n        reshape_prefers_feature_last_for_adjx_batch_matmul_fn=_reshape_prefers_feature_last_for_adjx_batch_matmul,\n    )\n    _tensor_exact_static_shape_list = lambda tensor_name: _tensor_exact_static_shape_list_for_model_ir(\n        model_ir=model_ir,\n        tensor_name=tensor_name,\n    )\n    _static_sequence_length = lambda tensor_name: _static_sequence_length_for_model_ir(\n        model_ir=model_ir,\n        tensor_name=tensor_name,\n    )\n    _is_identity_nms_postprocess_gather = lambda params_name, indices_name: _is_identity_nms_postprocess_gather_for_codegen(\n        model_ir=model_ir,\n        tensor_expr_aliases=tensor_expr_aliases,\n        producer_index=producer_index,\n        scalar_literal_expr_fn=_scalar_literal_expr,\n        params_name=params_name,\n        indices_name=indices_name,\n    )\n    _range_only_feeds_identity_nms_postprocess_gathers = lambda output_name: _range_only_feeds_identity_nms_postprocess_gathers_for_codegen(\n        model_ir=model_ir,\n        consumer_index=consumer_index,\n        is_identity_nms_postprocess_gather_fn=_is_identity_nms_postprocess_gather,\n        output_name=output_name,\n    )\n    _reshape_preserves_channel_last_sequence = lambda input_shape, output_shape, input_layout: _reshape_preserves_channel_last_sequence_for_codegen(\n        input_shape=input_shape,\n        output_shape=output_shape,\n        input_layout=input_layout,\n    )\n    _reshape_prefers_feature_last_for_adjx_batch_matmul = lambda input_tensor_name, output_name: _reshape_prefers_feature_last_for_adjx_batch_matmul_for_codegen(\n        model_ir=model_ir,\n        consumer_index=consumer_index,\n        input_tensor_name=input_tensor_name,\n        output_name=output_name,\n    )\n    _matmul_broadcast_shape = lambda lhs_batch, rhs_batch: _matmul_broadcast_shape_for_codegen(\n        lhs_batch=lhs_batch,\n        rhs_batch=rhs_batch,\n    )\n    _infer_batch_matmul_shape = lambda lhs_shape, rhs_shape, *, adj_x, adj_y: _infer_batch_matmul_shape_for_codegen(\n        lhs_shape=lhs_shape,\n        rhs_shape=rhs_shape,\n        adj_x=adj_x,\n        adj_y=adj_y,\n    )\n    _infer_reduction_shape = lambda input_shape, axes, *, keepdims: _infer_reduction_shape_for_codegen(\n        input_shape=input_shape,\n        axes=axes,\n        keepdims=keepdims,\n    )\n    _infer_gather_nd_shape = lambda params_shape, indices_tensor_name: _infer_gather_nd_shape_for_codegen(\n        model_ir=model_ir,\n        params_shape=params_shape,\n        indices_tensor_name=indices_tensor_name,\n    )\n    _emit_maybe_aligned_expr = lambda *, output_name, expr, inferred_shape: _emit_maybe_aligned_expr_for_codegen(\n        runtime_imports=runtime_imports,\n        output_name=output_name,\n        expr=expr,\n        inferred_shape=inferred_shape,\n        tensor_shape_list_fn=_tensor_shape_list,\n        target_shape_literal_fn=_target_shape_literal,\n    )\n    _emit_module_output_expr = lambda *, output_name, expr, raw_output_layout: _emit_module_output_expr_for_codegen(\n        model_ir=model_ir,\n        runtime_imports=runtime_imports,\n        output_name=output_name,\n        expr=expr,\n        raw_output_layout=raw_output_layout,\n        tensor_shape_list_fn=_tensor_shape_list,\n        target_shape_literal_fn=_target_shape_literal,\n    )\n    _is_constant_tensor_name = lambda tensor_name: _is_constant_tensor_name_for_codegen(\n        model_ir=model_ir,\n        tensor_name=tensor_name,\n    )\n    _shape_tensor_constant_is_non_zero_int_vector = lambda tensor_name: _shape_tensor_constant_is_non_zero_int_vector_for_codegen(\n        model_ir=model_ir,\n        tensor_name=tensor_name,\n    )\n    _static_int_tensor_values = lambda tensor_name, *, _visited=None: _static_int_tensor_values_for_codegen(\n        model_ir=model_ir,\n        producer_index=producer_index,\n        tensor_name=tensor_name,\n        _visited=_visited,\n    )\n    _reshape_shape_tensor_uses_runtime_dims = lambda tensor_name, *, _visited=None: _reshape_shape_tensor_uses_runtime_dims_for_codegen(\n        model_ir=model_ir,\n        producer_index=producer_index,\n        tensor_name=tensor_name,\n        _visited=_visited,\n    )\n    _should_skip_align_for_shape_preserving_unary = lambda input_name, output_name: _should_skip_align_for_shape_preserving_unary_for_codegen(\n        model_ir=model_ir,\n        input_name=input_name,\n        output_name=output_name,\n        tensor_shape_list_fn=_tensor_shape_list,\n    )\n    _next_unique_attr_name = lambda base_name: _next_unique_attr_name_for_codegen(\n        base_name=base_name,\n        module_attr_counts=module_attr_counts,\n        affine_layer_norm_specs=affine_layer_norm_specs,\n        op_module_attr_names=op_module_attr_names,\n    )\n    _canonical_codegen_name = _canonical_codegen_name_for_codegen\n    _match_affine_layer_norm = lambda op_index, op: _match_affine_layer_norm_for_codegen(\n        model_ir=model_ir,\n        producer_index=producer_index,\n        is_constant_tensor_name_fn=_is_constant_tensor_name,\n        canonical_codegen_name_fn=_canonical_codegen_name,\n        next_unique_attr_name_fn=_next_unique_attr_name,\n        op_index=op_index,\n        op=op,\n    )\n    _match_swish_activation_pattern = lambda tensor_name, consumer_indices: _match_swish_activation_pattern_for_codegen(\n        model_ir=model_ir,\n        consumer_index=consumer_index,\n        tensor_name=tensor_name,\n        consumer_indices=consumer_indices,\n    )\n    _topk_codegen_layout_bridge = lambda *, input_name, value_output_name, index_output_name: _topk_codegen_layout_bridge_for_codegen(\n        tensor_shape_list_fn=_tensor_shape_list,\n        input_name=input_name,\n        value_output_name=value_output_name,\n        index_output_name=index_output_name,\n    )\n    _pad_literal_expr = lambda tensor_name: _pad_literal_expr_for_codegen(\n        model_ir=model_ir,\n        tensor_name=tensor_name,\n    )\n    _constant_pad_pairs = lambda tensor_name: _constant_pad_pairs_for_codegen(\n        model_ir=model_ir,\n        tensor_name=tensor_name,\n    )\n    _scalar_literal_expr = lambda tensor_name: _scalar_literal_expr_for_codegen(\n        model_ir=model_ir,\n        tensor_name=tensor_name,\n    )\n    _int_scalar_literal_expr = lambda tensor_name: _int_scalar_literal_expr_for_codegen(\n        static_int_tensor_values_fn=lambda name: _static_int_tensor_values(name),\n        tensor_name=tensor_name,\n    )\n    _axis_expr_from_input = lambda tensor_name, *, device_expr: _axis_expr_from_input_for_codegen(\n        runtime_imports=runtime_imports,\n        int_scalar_literal_expr_fn=_int_scalar_literal_expr,\n        tensor_expr_fn=_tensor_expr,\n        tensor_name=tensor_name,\n        device_expr=device_expr,\n    )\n    _activation_lines = _activation_lines_for_codegen\n    _static_mirror_pad_expr = lambda *, input_tensor_name, pads_tensor_name: _static_mirror_pad_expr_for_codegen(\n        model_ir=model_ir,\n        runtime_imports=runtime_imports,\n        constant_pad_pairs_fn=_constant_pad_pairs,\n        tensor_expr_fn=_tensor_expr,\n        input_tensor_name=input_tensor_name,\n        pads_tensor_name=pads_tensor_name,\n    )\n    _is_sequential_single_input_graph = lambda: _is_sequential_single_input_graph_for_codegen(\n        model_ir=model_ir,\n    )\n    _is_channel_last_layout = _is_channel_last_layout_for_codegen\n    _tensor_dtype_name = lambda tensor_name: _tensor_dtype_name_for_codegen(\n        model_ir=model_ir,\n        tensor_name=tensor_name,\n    )\n    _can_emit_direct_module_call = lambda op: _can_emit_direct_module_call_for_codegen(\n        model_ir=model_ir,\n        is_channel_last_layout_fn=_is_channel_last_layout,\n        op=op,\n    )\n    _gather_input_pre_permute = lambda *, params_name, output_name, axis, batch_dims: _gather_input_pre_permute_for_codegen(\n        model_ir=model_ir,\n        params_name=params_name,\n        output_name=output_name,\n        axis=axis,\n        batch_dims=batch_dims,\n    )\n    _conv2d_output_spatial_shape = lambda *, input_hw, kernel_hw, stride_hw, dilation_hw, padding_mode: _conv2d_output_spatial_shape_for_codegen(\n        input_hw=input_hw,\n        kernel_hw=kernel_hw,\n        stride_hw=stride_hw,\n        dilation_hw=dilation_hw,\n        padding_mode=padding_mode,\n    )\n    _conv3d_output_spatial_shape = lambda *, input_dhw, kernel_dhw, stride_dhw, dilation_dhw, padding_mode: _conv3d_output_spatial_shape_for_codegen(\n        input_dhw=input_dhw,\n        kernel_dhw=kernel_dhw,\n        stride_dhw=stride_dhw,\n        dilation_dhw=dilation_dhw,\n        padding_mode=padding_mode,\n    )\n    _conv3d_transpose_output_spatial_shape = lambda *, input_dhw, kernel_dhw, stride_dhw, dilation_dhw, padding_mode: _conv3d_transpose_output_spatial_shape_for_codegen(\n        input_dhw=input_dhw,\n        kernel_dhw=kernel_dhw,\n        stride_dhw=stride_dhw,\n        dilation_dhw=dilation_dhw,\n        padding_mode=padding_mode,\n    )\n    _conv2d_same_pad_arg = lambda *, input_shape, output_shape, weight_shape, options, input_pre_permute=None, input_logical_layout=None, output_logical_layout=None: _conv2d_same_pad_padding_arg_for_codegen(\n        input_shape=input_shape,\n        output_shape=output_shape,\n        weight_shape=weight_shape,\n        options=options,\n        input_pre_permute=input_pre_permute,\n        input_logical_layout=input_logical_layout,\n        output_logical_layout=output_logical_layout,\n    )\n    _infer_conv2d_layout_candidate = lambda *, input_shape, output_shape, weight_shape, options, depthwise: _infer_conv2d_layout_candidate_for_codegen(\n        input_shape=input_shape,\n        output_shape=output_shape,\n        weight_shape=weight_shape,\n        options=options,\n        depthwise=depthwise,\n    )\n    _conv2d_input_pre_permute = lambda input_shape, output_shape, weight_shape, options, *, input_logical_layout=None, output_logical_layout=None, depthwise=False: _conv2d_input_pre_permute_for_codegen(\n        input_shape=input_shape,\n        output_shape=output_shape,\n        weight_shape=weight_shape,\n        options=options,\n        input_logical_layout=input_logical_layout,\n        output_logical_layout=output_logical_layout,\n        depthwise=depthwise,\n    )\n    _infer_effective_rank4_runtime_layout = lambda tensor_name: _infer_effective_rank4_runtime_layout_for_codegen(\n        model_ir=model_ir,\n        producer_index=producer_index,\n        consumer_index=consumer_index,\n        tensor_name=tensor_name,\n    )\n    _infer_conv2d_ctor_params = lambda *, input_shape, output_shape, weight_shape, options, input_logical_layout=None, output_logical_layout=None, depthwise: _infer_conv2d_ctor_params_for_codegen(\n        input_shape=input_shape,\n        output_shape=output_shape,\n        weight_shape=weight_shape,\n        options=options,\n        input_logical_layout=input_logical_layout,\n        output_logical_layout=output_logical_layout,\n        depthwise=depthwise,\n    )\n    _infer_conv3d_ctor_params = lambda *, input_shape, output_shape, weight_shape, options, input_logical_layout=None, output_logical_layout=None: _infer_conv3d_ctor_params_for_codegen(\n        input_shape=input_shape,\n        output_shape=output_shape,\n        weight_shape=weight_shape,\n        options=options,\n        input_logical_layout=input_logical_layout,\n        output_logical_layout=output_logical_layout,\n    )\n    _infer_conv3d_transpose_ctor_params = lambda *, input_shape, output_shape, weight_shape, options, input_logical_layout=None, output_logical_layout=None: _infer_conv3d_transpose_ctor_params_for_codegen(\n        input_shape=input_shape,\n        output_shape=output_shape,\n        weight_shape=weight_shape,\n        options=options,\n        input_logical_layout=input_logical_layout,\n        output_logical_layout=output_logical_layout,\n    )\n\n    for op_index, op in enumerate(model_ir.operators):\n        affine_layer_norm_spec = _match_affine_layer_norm(int(op_index), op)\n        if affine_layer_norm_spec is None:\n            continue\n        attr_name = str(affine_layer_norm_spec["attr_name"])\n        module_init_lines.extend(\n            [\n                f"self.{attr_name} = _AffineLayerNorm(",\n                f"    shape={repr(list(affine_layer_norm_spec[\'gamma_shape\']))},",\n                f"    dtype={_torch_dtype_literal(str(affine_layer_norm_spec[\'gamma_dtype\']))},",\n                ")",\n            ]\n        )\n        load_specs.append((f"{attr_name}.gamma", str(affine_layer_norm_spec["gamma_name"])))\n        load_specs.append((f"{attr_name}.beta", str(affine_layer_norm_spec["beta_name"])))\n        submodule_state_tensor_names.update(\n            {\n                str(affine_layer_norm_spec["gamma_name"]),\n                str(affine_layer_norm_spec["beta_name"]),\n            }\n        )\n        affine_layer_norm_specs[int(op_index)] = affine_layer_norm_spec\n        skipped_op_indices.add(int(affine_layer_norm_spec["mul_op_index"]))\n\n    for op_index, op in enumerate(model_ir.operators):\n        op_type = str(op.op_type)\n        if op_type not in _DIRECT_CODEGEN_MODULE_OP_TYPES:\n            continue\n        attr_base = _direct_codegen_module_attr_base(op_type)\n        attr_index = int(module_attr_counts.get(attr_base, 0))\n        module_attr_counts[attr_base] = attr_index + 1\n        attr_name = f"{attr_base}_{attr_index}"\n        op_module_attr_names[int(op_index)] = attr_name\n        if op_type == "UNIDIRECTIONAL_SEQUENCE_RNN":\n            if not _can_direct_codegen_sequence_rnn_op(model_ir, op):\n                raise ModelIRPyTorchExportError(\n                    "Native PyTorch-like model.py codegen does not support this UNIDIRECTIONAL_SEQUENCE_RNN configuration."\n                )\n            weight_name = _sequence_lstm_input_name(op, 1)\n            recurrent_name = _sequence_lstm_input_name(op, 2)\n            bias_name = _sequence_lstm_input_name(op, 3)\n            weight = _require_constant_array_from_model_ir(\n                model_ir=model_ir,\n                tensor_name=weight_name,\n                context="RNN input weight",\n            )\n            hidden_size = int(weight.shape[0])\n            input_size = int(weight.shape[1])\n            weight_dtype = str(model_ir.tensors[weight_name].dtype)\n            bias_hh_name = _add_synthetic_tensor_to_model_ir(model_ir=model_ir, synthetic_tensor_serial_ref=synthetic_tensor_serial_ref, \n                base_name=f"{attr_name}_bias_hh_l0",\n                data=np.zeros((hidden_size,), dtype=np.float32),\n                dtype=weight_dtype,\n            )\n            module_param_tensor_names.update(\n                {\n                    str(weight_name),\n                    str(recurrent_name),\n                    str(bias_name),\n                    str(bias_hh_name),\n                }\n            )\n            activation_name = str(op.options.get("fusedActivationFunction", "TANH")).strip().lower()\n            module_init_lines.extend(\n                [\n                    f"self.{attr_name} = _SequenceRNNBlock(",\n                    f"    input_size={input_size},",\n                    f"    hidden_size={hidden_size},",\n                    f"    activation={activation_name!r},",\n                    ")",\n                ]\n            )\n            load_specs.extend(\n                [\n                    (f"{attr_name}.rnn.weight_ih_l0", str(weight_name)),\n                    (f"{attr_name}.rnn.weight_hh_l0", str(recurrent_name)),\n                    (f"{attr_name}.rnn.bias_ih_l0", str(bias_name)),\n                    (f"{attr_name}.rnn.bias_hh_l0", str(bias_hh_name)),\n                ]\n            )\n            continue\n        if op_type == "UNIDIRECTIONAL_SEQUENCE_LSTM":\n            if not _can_direct_codegen_sequence_lstm_op(model_ir, op):\n                raise ModelIRPyTorchExportError(\n                    "Native PyTorch-like model.py codegen does not support this UNIDIRECTIONAL_SEQUENCE_LSTM configuration."\n                )\n            index_spec = _sequence_lstm_index_spec(op)\n            if index_spec is None:\n                raise ModelIRPyTorchExportError(\n                    "Native PyTorch-like model.py codegen could not resolve UNIDIRECTIONAL_SEQUENCE_LSTM input layout."\n                )\n            weight_input_names = [\n                _sequence_lstm_input_name(op, idx)\n                for idx in list(index_spec["weight_input_indices"])\n            ]\n            recurrent_input_names = [\n                _sequence_lstm_input_name(op, idx)\n                for idx in list(index_spec["recurrent_input_indices"])\n            ]\n            input_weight = _require_constant_array_from_model_ir(\n                model_ir=model_ir,\n                tensor_name=weight_input_names[0],\n                context="LSTM input weight",\n            )\n            hidden_size = int(input_weight.shape[0])\n            input_size = int(input_weight.shape[1])\n            weight_dtype = str(model_ir.tensors[weight_input_names[0]].dtype)\n            bias_ih_name = _sequence_lstm_bias_array_for_model_ir(model_ir=model_ir, synthetic_tensor_serial_ref=synthetic_tensor_serial_ref, \n                op=op,\n                indices=list(index_spec["bias_indices"]),\n                hidden_size=hidden_size,\n                dtype=weight_dtype,\n                base_name=f"{attr_name}_bias_ih_l0",\n            )\n            bias_hh_name = _add_synthetic_tensor_to_model_ir(model_ir=model_ir, synthetic_tensor_serial_ref=synthetic_tensor_serial_ref, \n                base_name=f"{attr_name}_bias_hh_l0",\n                data=np.zeros((4 * hidden_size,), dtype=np.float32),\n                dtype=weight_dtype,\n            )\n            weight_ih_name = _add_synthetic_tensor_to_model_ir(model_ir=model_ir, synthetic_tensor_serial_ref=synthetic_tensor_serial_ref, \n                base_name=f"{attr_name}_weight_ih_l0",\n                data=np.concatenate(\n                    [\n                        _require_constant_array_from_model_ir(\n                            model_ir=model_ir,\n                            tensor_name=name,\n                            context=f"LSTM input gate weight {gate_name}",\n                        )\n                        for gate_name, name in zip(["i", "f", "c", "o"], weight_input_names)\n                    ],\n                    axis=0,\n                ).astype(np.float32, copy=False),\n                dtype=weight_dtype,\n            )\n            weight_hh_name = _add_synthetic_tensor_to_model_ir(model_ir=model_ir, synthetic_tensor_serial_ref=synthetic_tensor_serial_ref, \n                base_name=f"{attr_name}_weight_hh_l0",\n                data=np.concatenate(\n                    [\n                        _require_constant_array_from_model_ir(\n                            model_ir=model_ir,\n                            tensor_name=name,\n                            context=f"LSTM recurrent gate weight {gate_name}",\n                        )\n                        for gate_name, name in zip(["i", "f", "c", "o"], recurrent_input_names)\n                    ],\n                    axis=0,\n                ).astype(np.float32, copy=False),\n                dtype=weight_dtype,\n            )\n            module_param_tensor_names.update(\n                {\n                    str(weight_ih_name),\n                    str(weight_hh_name),\n                    str(bias_ih_name),\n                    str(bias_hh_name),\n                }\n            )\n            module_init_lines.extend(\n                [\n                    f"self.{attr_name} = _SequenceLSTMBlock(",\n                    f"    input_size={input_size},",\n                    f"    hidden_size={hidden_size},",\n                    f"    sequence_length={_static_sequence_length(str(op.inputs[0]))},",\n                    "    bidirectional=False,",\n                    "    merge_outputs=True,",\n                    ")",\n                ]\n            )\n            load_specs.extend(\n                [\n                    (f"{attr_name}.lstm.weight_ih_l0", str(weight_ih_name)),\n                    (f"{attr_name}.lstm.weight_hh_l0", str(weight_hh_name)),\n                    (f"{attr_name}.lstm.bias_ih_l0", str(bias_ih_name)),\n                    (f"{attr_name}.lstm.bias_hh_l0", str(bias_hh_name)),\n                ]\n            )\n            continue\n        if op_type == "BIDIRECTIONAL_SEQUENCE_LSTM":\n            if not _can_direct_codegen_sequence_lstm_op(model_ir, op):\n                raise ModelIRPyTorchExportError(\n                    "Native PyTorch-like model.py codegen does not support this BIDIRECTIONAL_SEQUENCE_LSTM configuration."\n                )\n            index_spec = _sequence_lstm_index_spec(op)\n            if index_spec is None:\n                raise ModelIRPyTorchExportError(\n                    "Native PyTorch-like model.py codegen could not resolve BIDIRECTIONAL_SEQUENCE_LSTM input layout."\n                )\n            fw_weight_input_names = [\n                _sequence_lstm_input_name(op, idx)\n                for idx in list(index_spec["fw_weight_input_indices"])\n            ]\n            fw_recurrent_input_names = [\n                _sequence_lstm_input_name(op, idx)\n                for idx in list(index_spec["fw_recurrent_input_indices"])\n            ]\n            bw_weight_input_names = [\n                _sequence_lstm_input_name(op, idx)\n                for idx in list(index_spec["bw_weight_input_indices"])\n            ]\n            bw_recurrent_input_names = [\n                _sequence_lstm_input_name(op, idx)\n                for idx in list(index_spec["bw_recurrent_input_indices"])\n            ]\n            fw_input_weight = _require_constant_array_from_model_ir(\n                model_ir=model_ir,\n                tensor_name=fw_weight_input_names[0],\n                context="forward LSTM input weight",\n            )\n            hidden_size = int(fw_input_weight.shape[0])\n            input_size = int(fw_input_weight.shape[1])\n            weight_dtype = str(model_ir.tensors[fw_weight_input_names[0]].dtype)\n            fw_bias_ih_name = _sequence_lstm_bias_array_for_model_ir(model_ir=model_ir, synthetic_tensor_serial_ref=synthetic_tensor_serial_ref, \n                op=op,\n                indices=list(index_spec["fw_bias_indices"]),\n                hidden_size=hidden_size,\n                dtype=weight_dtype,\n                base_name=f"{attr_name}_bias_ih_l0",\n            )\n            fw_bias_hh_name = _add_synthetic_tensor_to_model_ir(model_ir=model_ir, synthetic_tensor_serial_ref=synthetic_tensor_serial_ref, \n                base_name=f"{attr_name}_bias_hh_l0",\n                data=np.zeros((4 * hidden_size,), dtype=np.float32),\n                dtype=weight_dtype,\n            )\n            bw_bias_ih_name = _sequence_lstm_bias_array_for_model_ir(model_ir=model_ir, synthetic_tensor_serial_ref=synthetic_tensor_serial_ref, \n                op=op,\n                indices=list(index_spec["bw_bias_indices"]),\n                hidden_size=hidden_size,\n                dtype=weight_dtype,\n                base_name=f"{attr_name}_bias_ih_l0_reverse",\n            )\n            bw_bias_hh_name = _add_synthetic_tensor_to_model_ir(model_ir=model_ir, synthetic_tensor_serial_ref=synthetic_tensor_serial_ref, \n                base_name=f"{attr_name}_bias_hh_l0_reverse",\n                data=np.zeros((4 * hidden_size,), dtype=np.float32),\n                dtype=weight_dtype,\n            )\n            fw_weight_ih_name = _add_synthetic_tensor_to_model_ir(model_ir=model_ir, synthetic_tensor_serial_ref=synthetic_tensor_serial_ref, \n                base_name=f"{attr_name}_weight_ih_l0",\n                data=np.concatenate(\n                    [\n                        _require_constant_array_from_model_ir(\n                            model_ir=model_ir,\n                            tensor_name=name,\n                            context=f"forward LSTM input gate weight {gate_name}",\n                        )\n                        for gate_name, name in zip(["i", "f", "c", "o"], fw_weight_input_names)\n                    ],\n                    axis=0,\n                ).astype(np.float32, copy=False),\n                dtype=weight_dtype,\n            )\n            fw_weight_hh_name = _add_synthetic_tensor_to_model_ir(model_ir=model_ir, synthetic_tensor_serial_ref=synthetic_tensor_serial_ref, \n                base_name=f"{attr_name}_weight_hh_l0",\n                data=np.concatenate(\n                    [\n                        _require_constant_array_from_model_ir(\n                            model_ir=model_ir,\n                            tensor_name=name,\n                            context=f"forward LSTM recurrent gate weight {gate_name}",\n                        )\n                        for gate_name, name in zip(["i", "f", "c", "o"], fw_recurrent_input_names)\n                    ],\n                    axis=0,\n                ).astype(np.float32, copy=False),\n                dtype=weight_dtype,\n            )\n            bw_weight_ih_name = _add_synthetic_tensor_to_model_ir(model_ir=model_ir, synthetic_tensor_serial_ref=synthetic_tensor_serial_ref, \n                base_name=f"{attr_name}_weight_ih_l0_reverse",\n                data=np.concatenate(\n                    [\n                        _require_constant_array_from_model_ir(\n                            model_ir=model_ir,\n                            tensor_name=name,\n                            context=f"backward LSTM input gate weight {gate_name}",\n                        )\n                        for gate_name, name in zip(["i", "f", "c", "o"], bw_weight_input_names)\n                    ],\n                    axis=0,\n                ).astype(np.float32, copy=False),\n                dtype=weight_dtype,\n            )\n            bw_weight_hh_name = _add_synthetic_tensor_to_model_ir(model_ir=model_ir, synthetic_tensor_serial_ref=synthetic_tensor_serial_ref, \n                base_name=f"{attr_name}_weight_hh_l0_reverse",\n                data=np.concatenate(\n                    [\n                        _require_constant_array_from_model_ir(\n                            model_ir=model_ir,\n                            tensor_name=name,\n                            context=f"backward LSTM recurrent gate weight {gate_name}",\n                        )\n                        for gate_name, name in zip(["i", "f", "c", "o"], bw_recurrent_input_names)\n                    ],\n                    axis=0,\n                ).astype(np.float32, copy=False),\n                dtype=weight_dtype,\n            )\n            module_param_tensor_names.update(\n                {\n                    str(fw_weight_ih_name),\n                    str(fw_weight_hh_name),\n                    str(fw_bias_ih_name),\n                    str(fw_bias_hh_name),\n                    str(bw_weight_ih_name),\n                    str(bw_weight_hh_name),\n                    str(bw_bias_ih_name),\n                    str(bw_bias_hh_name),\n                }\n            )\n            module_init_lines.extend(\n                [\n                    f"self.{attr_name} = _SequenceLSTMBlock(",\n                    f"    input_size={input_size},",\n                    f"    hidden_size={hidden_size},",\n                    f"    sequence_length={_static_sequence_length(str(op.inputs[0]))},",\n                    "    bidirectional=True,",\n                    f"    merge_outputs={bool(op.options.get(\'mergeOutputs\', True))},",\n                    ")",\n                ]\n            )\n            load_specs.extend(\n                [\n                    (f"{attr_name}.lstm.weight_ih_l0", str(fw_weight_ih_name)),\n                    (f"{attr_name}.lstm.weight_hh_l0", str(fw_weight_hh_name)),\n                    (f"{attr_name}.lstm.bias_ih_l0", str(fw_bias_ih_name)),\n                    (f"{attr_name}.lstm.bias_hh_l0", str(fw_bias_hh_name)),\n                    (f"{attr_name}.lstm.weight_ih_l0_reverse", str(bw_weight_ih_name)),\n                    (f"{attr_name}.lstm.weight_hh_l0_reverse", str(bw_weight_hh_name)),\n                    (f"{attr_name}.lstm.bias_ih_l0_reverse", str(bw_bias_ih_name)),\n                    (f"{attr_name}.lstm.bias_hh_l0_reverse", str(bw_bias_hh_name)),\n                ]\n            )\n            continue\n        weight_name = str(op.inputs[1]) if len(op.inputs) >= 2 else ""\n        bias_name = ""\n        if op_type in {"CONV_2D", "DEPTHWISE_CONV_2D", "CONV_3D"} and len(op.inputs) >= 3 and str(op.inputs[2]) != "":\n            bias_name = str(op.inputs[2])\n        elif op_type == "FULLY_CONNECTED" and len(op.inputs) >= 3 and str(op.inputs[2]) != "":\n            bias_name = str(op.inputs[2])\n        elif op_type in {"TRANSPOSE_CONV", "CONV_3D_TRANSPOSE"} and len(op.inputs) >= 4 and str(op.inputs[3]) != "":\n            bias_name = str(op.inputs[3])\n        module_param_tensor_names.add(weight_name)\n        if bias_name != "":\n            module_param_tensor_names.add(bias_name)\n        weight_tensor = model_ir.tensors[weight_name]\n        input_tensor_name = str(op.inputs[0]) if op_type not in {"TRANSPOSE_CONV", "CONV_3D_TRANSPOSE"} else str(op.inputs[2])\n        input_tensor = model_ir.tensors[input_tensor_name]\n        options = dict(op.options)\n        conv_input_pre_permute: Optional[List[int]] = None\n        conv_same_pad_arg: Optional[List[int]] = None\n        if op_type in {"CONV_2D", "DEPTHWISE_CONV_2D"}:\n            conv_input_pre_permute = _conv2d_input_pre_permute(\n                _tensor_shape_list(str(input_tensor_name)),\n                _tensor_shape_list(str(op.outputs[0])),\n                _tensor_shape_list(str(op.inputs[1])),\n                options,\n                input_logical_layout=normalize_logical_layout(input_tensor.logical_layout),\n                output_logical_layout=normalize_logical_layout(\n                    model_ir.tensors[str(op.outputs[0])].logical_layout\n                ),\n                depthwise=(op_type == "DEPTHWISE_CONV_2D"),\n            )\n            conv_same_pad_arg = _conv2d_same_pad_arg(\n                input_shape=_tensor_shape_list(str(input_tensor_name)),\n                output_shape=_tensor_shape_list(str(op.outputs[0])),\n                weight_shape=_tensor_shape_list(str(op.inputs[1])),\n                options=options,\n                input_pre_permute=conv_input_pre_permute,\n                input_logical_layout=normalize_logical_layout(input_tensor.logical_layout),\n                output_logical_layout=normalize_logical_layout(\n                    model_ir.tensors[str(op.outputs[0])].logical_layout\n                ),\n            )\n            conv_module_pad_specs[int(op_index)] = conv_same_pad_arg\n        fused_block_input_name = str(input_tensor_name)\n        fused_block_output_name = str(op.outputs[0]) if len(op.outputs) == 1 else ""\n        fused_block_pad: Optional[str] = None\n        fused_block_pad_mode = "constant"\n        fused_block_pad_value: Optional[str] = None\n        fused_bridge_output_op_idx: Optional[int] = None\n        fused_block_activation, fused_block_negative_slope = _conv_block_activation_config_from_fused_name(\n            str(op.options.get("fusedActivationFunction", "NONE")),\n            alpha=(\n                float(op.options.get("alpha", 0.2))\n                if "alpha" in op.options\n                else None\n            ),\n        )\n        if op_type in {"CONV_2D", "DEPTHWISE_CONV_2D"} and len(op.outputs) == 1:\n            (\n                fused_block_input_name,\n                conv_input_pre_permute,\n                folded_public_bridge_op_idx,\n            ) = _fold_single_consumer_public_input_bridge(\n                tensor_name=str(fused_block_input_name),\n                downstream_permute=conv_input_pre_permute,\n            )\n            if folded_public_bridge_op_idx is not None:\n                skipped_op_indices.add(int(folded_public_bridge_op_idx))\n            input_producer_idx = producer_index.get(str(fused_block_input_name), None)\n            if input_producer_idx is not None:\n                input_producer = model_ir.operators[int(input_producer_idx)]\n                input_producer_type = str(input_producer.op_type)\n                if (\n                    input_producer_type in {"PAD", "PADV2"}\n                    and len(consumer_index.get(str(fused_block_input_name), [])) == 1\n                    and len(input_producer.inputs) >= 2\n                ):\n                    pad_axis_permutation = (\n                        [int(v) for v in list(conv_input_pre_permute)]\n                        if conv_input_pre_permute is not None\n                        else None\n                    )\n                    if pad_axis_permutation is None:\n                        pad_input_tensor = model_ir.tensors.get(str(input_producer.inputs[0]), None)\n                        pad_output_tensor = model_ir.tensors.get(str(input_producer.outputs[0]), None)\n                        if pad_input_tensor is not None and pad_output_tensor is not None:\n                            pad_input_layout = normalize_logical_layout(pad_input_tensor.logical_layout)\n                            pad_output_layout = normalize_logical_layout(pad_output_tensor.logical_layout)\n                            pad_rank = len(list(pad_output_tensor.shape))\n                            inferred_pad_perm = logical_layout_permutation(\n                                source_layout=pad_input_layout,\n                                target_layout=pad_output_layout,\n                            )\n                            if inferred_pad_perm is not None and len(inferred_pad_perm) == pad_rank:\n                                pad_axis_permutation = [int(v) for v in list(inferred_pad_perm)]\n                    original_conv_input_pre_permute = (\n                        [int(v) for v in list(conv_input_pre_permute)]\n                        if conv_input_pre_permute is not None\n                        else None\n                    )\n                    pad_literal = _torch_pad_literal_for_constant_tensor(\n                        model_ir.tensors.get(str(input_producer.inputs[1]), None),\n                        axis_permutation=pad_axis_permutation,\n                    )\n                    if pad_literal is not None:\n                        fused_block_input_name = str(input_producer.inputs[0])\n                        fused_block_pad = pad_literal\n                        pad_matches_pre_permuted_input = _pad_output_matches_pre_permuted_input(\n                            input_tensor=model_ir.tensors.get(str(fused_block_input_name), None),\n                            output_tensor=model_ir.tensors.get(str(input_producer.outputs[0]), None),\n                            pads_tensor=model_ir.tensors.get(str(input_producer.inputs[1]), None),\n                            input_pre_permute=conv_input_pre_permute,\n                        )\n                        if pad_matches_pre_permuted_input:\n                            rewrite_pad_axis_permutation = original_conv_input_pre_permute\n                            if rewrite_pad_axis_permutation is None:\n                                pad_input_tensor = model_ir.tensors.get(str(input_producer.inputs[0]), None)\n                                pad_output_tensor = model_ir.tensors.get(str(input_producer.outputs[0]), None)\n                                if pad_input_tensor is not None and pad_output_tensor is not None:\n                                    pad_input_layout = normalize_logical_layout(pad_input_tensor.logical_layout)\n                                    pad_output_layout = normalize_logical_layout(pad_output_tensor.logical_layout)\n                                    pad_rank = len(list(pad_output_tensor.shape))\n                                    inferred_pad_perm = logical_layout_permutation(\n                                        source_layout=pad_output_layout,\n                                        target_layout=pad_input_layout,\n                                    )\n                                    if inferred_pad_perm is not None and len(inferred_pad_perm) == pad_rank:\n                                        rewrite_pad_axis_permutation = [int(v) for v in list(inferred_pad_perm)]\n                            if rewrite_pad_axis_permutation is not None:\n                                rewritten_pad_literal = _torch_pad_literal_for_constant_tensor(\n                                    model_ir.tensors.get(str(input_producer.inputs[1]), None),\n                                    axis_permutation=rewrite_pad_axis_permutation,\n                                )\n                                if rewritten_pad_literal is not None:\n                                    pad_literal = rewritten_pad_literal\n                                    fused_block_pad = pad_literal\n                            conv_input_pre_permute = None\n                        upstream_producer_idx = producer_index.get(str(fused_block_input_name), None)\n                        if upstream_producer_idx is not None:\n                            upstream_producer = model_ir.operators[int(upstream_producer_idx)]\n                            if (\n                                str(upstream_producer.op_type) == "TRANSPOSE"\n                                and len(upstream_producer.outputs) == 1\n                                and len(consumer_index.get(str(fused_block_input_name), [])) == 1\n                            ):\n                                upstream_perm = _read_transpose_perm(model_ir, upstream_producer)\n                                composed_perm = (\n                                    None\n                                    if pad_matches_pre_permuted_input\n                                    else _compose_axis_permutations(\n                                        upstream_perm,\n                                        conv_input_pre_permute,\n                                    )\n                                )\n                                if upstream_perm is not None or conv_input_pre_permute is not None:\n                                    fused_block_input_name = str(upstream_producer.inputs[0])\n                                    conv_input_pre_permute = composed_perm\n                                    skipped_op_indices.add(int(upstream_producer_idx))\n                        if input_producer_type == "PADV2" and len(input_producer.inputs) >= 3:\n                            scalar_literal = _scalar_literal_for_constant_tensor(\n                                model_ir.tensors.get(str(input_producer.inputs[2]), None)\n                            )\n                            if scalar_literal is not None:\n                                fused_block_pad_value = scalar_literal\n                        skipped_op_indices.add(int(input_producer_idx))\n                        inlined_constant_tensor_names.add(str(input_producer.inputs[1]))\n                        if input_producer_type == "PADV2" and len(input_producer.inputs) >= 3:\n                            inlined_constant_tensor_names.add(str(input_producer.inputs[2]))\n            bridged_output = _match_single_consumer_layout_bridge_transpose(\n                tensor_name=str(op.outputs[0]),\n                required_output_layout=channel_first_logical_layout(\n                    len(list(model_ir.tensors[str(op.outputs[0])].shape))\n                ),\n            )\n            effective_output_name = str(op.outputs[0])\n            output_consumer_indices = consumer_index.get(str(op.outputs[0]), [])\n            if bridged_output is not None:\n                candidate_output_name, candidate_bridge_output_op_idx = bridged_output\n                candidate_output_consumer_indices = consumer_index.get(str(candidate_output_name), [])\n                can_fold_bridge_output = False\n                candidate_swish_match = _match_swish_activation_pattern(\n                    str(candidate_output_name),\n                    candidate_output_consumer_indices,\n                )\n                if len(candidate_output_consumer_indices) == 1:\n                    candidate_activation_op = model_ir.operators[int(candidate_output_consumer_indices[0])]\n                    candidate_activation_type = str(candidate_activation_op.op_type)\n                    can_fold_bridge_output = (\n                        len(candidate_activation_op.inputs) == 1\n                        and len(candidate_activation_op.outputs) == 1\n                        and candidate_activation_type in {\n                            "LEAKY_RELU",\n                            "LOGISTIC",\n                            "RELU",\n                            "RELU6",\n                            "RELU_N1_TO_1",\n                            "RELU_0_TO_1",\n                            "TANH",\n                        }\n                    )\n                elif candidate_swish_match is not None and fused_block_activation == "none":\n                    can_fold_bridge_output = True\n                elif str(candidate_output_name) in {str(name) for name in list(model_ir.outputs)}:\n                    can_fold_bridge_output = True\n                if can_fold_bridge_output:\n                    effective_output_name = str(candidate_output_name)\n                    fused_bridge_output_op_idx = int(candidate_bridge_output_op_idx)\n                    output_consumer_indices = candidate_output_consumer_indices\n            swish_match = (\n                _match_swish_activation_pattern(\n                    str(effective_output_name),\n                    output_consumer_indices,\n                )\n                if fused_block_activation == "none"\n                else None\n            )\n            if swish_match is not None:\n                fused_block_activation = "silu"\n                fused_block_output_name = str(swish_match[0])\n                for matched_op_idx in sorted(int(idx) for idx in swish_match[1]):\n                    skipped_op_indices.add(int(matched_op_idx))\n                if fused_bridge_output_op_idx is not None:\n                    skipped_op_indices.add(int(fused_bridge_output_op_idx))\n            elif len(output_consumer_indices) == 1:\n                activation_op = model_ir.operators[int(output_consumer_indices[0])]\n                activation_type = str(activation_op.op_type)\n                if activation_type in {\n                    "LEAKY_RELU",\n                    "LOGISTIC",\n                    "RELU",\n                    "RELU6",\n                    "RELU_N1_TO_1",\n                    "RELU_0_TO_1",\n                    "TANH",\n                } and len(activation_op.inputs) == 1 and len(activation_op.outputs) == 1 and fused_block_activation == "none":\n                    fused_block_activation, fused_block_negative_slope = _conv_block_activation_config(activation_op)\n                    fused_block_output_name = str(activation_op.outputs[0])\n                    skipped_op_indices.add(int(output_consumer_indices[0]))\n                    if fused_bridge_output_op_idx is not None:\n                        skipped_op_indices.add(int(fused_bridge_output_op_idx))\n            elif fused_bridge_output_op_idx is not None:\n                fused_block_output_name = str(effective_output_name)\n                skipped_op_indices.add(int(fused_bridge_output_op_idx))\n        use_conv_block = op_type in {"CONV_2D", "DEPTHWISE_CONV_2D"}\n        if use_conv_block:\n            attr_base = "conv_block"\n            attr_index = int(module_attr_counts.get(attr_base, 0))\n            module_attr_counts[attr_base] = attr_index + 1\n            attr_name = f"{attr_base}_{attr_index}"\n            op_module_attr_names[int(op_index)] = attr_name\n            fused_module_specs[int(op_index)] = {\n                "input_name": str(fused_block_input_name),\n                "output_name": str(fused_block_output_name),\n                "input_pre_permute": conv_input_pre_permute,\n                "pad": fused_block_pad if fused_block_pad is not None else (repr(conv_same_pad_arg) if conv_same_pad_arg is not None else None),\n                "pad_mode": fused_block_pad_mode,\n                "pad_value": fused_block_pad_value,\n                "activation": fused_block_activation,\n                "negative_slope": fused_block_negative_slope,\n            }\n        if op_type == "CONV_2D":\n            conv_padding = (\n                [0, 0]\n                if conv_same_pad_arg is not None\n                else [\n                    int(((int(weight_tensor.shape[2]) - 1) * int(options.get(\'dilationHFactor\', 1))) // 2) if str(options.get(\'padding\', \'SAME\')).upper() != \'VALID\' else 0,\n                    int(((int(weight_tensor.shape[3]) - 1) * int(options.get(\'dilationWFactor\', 1))) // 2) if str(options.get(\'padding\', \'SAME\')).upper() != \'VALID\' else 0,\n                ]\n            )\n            conv_in_channels, conv_groups = _infer_conv2d_ctor_params(\n                input_shape=list(input_tensor.shape),\n                output_shape=list(model_ir.tensors[str(op.outputs[0])].shape),\n                weight_shape=list(weight_tensor.shape),\n                options=options,\n                input_logical_layout=input_tensor.logical_layout,\n                output_logical_layout=model_ir.tensors[str(op.outputs[0])].logical_layout,\n                depthwise=False,\n            )\n            conv_ctor_lines = [\n                "torch.nn.Conv2d(",\n                f"    in_channels={conv_in_channels},",\n                f"    out_channels={int(weight_tensor.shape[0])},",\n                f"    kernel_size={_shape_literal(weight_tensor.shape[2:])},",\n                f"    stride={_shape_literal([int(options.get(\'strideH\', 1)), int(options.get(\'strideW\', 1))])},",\n                f"    padding={_shape_literal(conv_padding)},",\n                f"    dilation={_shape_literal([int(options.get(\'dilationHFactor\', 1)), int(options.get(\'dilationWFactor\', 1))])},",\n                f"    groups={conv_groups},",\n                f"    bias={str(bias_name != \'\')},",\n                ")",\n            ]\n            if use_conv_block:\n                module_init_lines.extend(\n                    [\n                        f"self.{attr_name} = _Conv2dBlock(",\n                        f"    in_channels={conv_in_channels},",\n                        f"    out_channels={int(weight_tensor.shape[0])},",\n                        f"    kernel_size={_shape_literal(weight_tensor.shape[2:])},",\n                        f"    stride={_shape_literal([int(options.get(\'strideH\', 1)), int(options.get(\'strideW\', 1))])},",\n                        f"    padding={_shape_literal(conv_padding)},",\n                        f"    dilation={_shape_literal([int(options.get(\'dilationHFactor\', 1)), int(options.get(\'dilationWFactor\', 1))])},",\n                        f"    groups={conv_groups},",\n                        f"    bias={str(bias_name != \'\')},",\n                        f"    pad={fused_block_pad if fused_block_pad is not None else (repr(conv_same_pad_arg) if conv_same_pad_arg is not None else \'None\')},",\n                        f"    activation={fused_block_activation!r},",\n                        f"    negative_slope={repr(fused_block_negative_slope if fused_block_negative_slope is not None else 0.2)},",\n                        f"    pad_mode={fused_block_pad_mode!r},",\n                        f"    pad_value={fused_block_pad_value if fused_block_pad_value is not None else \'0.0\'},",\n                        ")",\n                    ]\n                )\n            else:\n                module_init_lines.extend(\n                    [\n                        f"self.{attr_name} = {conv_ctor_lines[0]}",\n                        *conv_ctor_lines[1:],\n                    ]\n                )\n        elif op_type == "DEPTHWISE_CONV_2D":\n            conv_padding = (\n                [0, 0]\n                if conv_same_pad_arg is not None\n                else [\n                    int(((int(weight_tensor.shape[2]) - 1) * int(options.get(\'dilationHFactor\', 1))) // 2) if str(options.get(\'padding\', \'SAME\')).upper() != \'VALID\' else 0,\n                    int(((int(weight_tensor.shape[3]) - 1) * int(options.get(\'dilationWFactor\', 1))) // 2) if str(options.get(\'padding\', \'SAME\')).upper() != \'VALID\' else 0,\n                ]\n            )\n            depthwise_in_channels, depthwise_groups = _infer_conv2d_ctor_params(\n                input_shape=list(input_tensor.shape),\n                output_shape=list(model_ir.tensors[str(op.outputs[0])].shape),\n                weight_shape=list(weight_tensor.shape),\n                options=options,\n                input_logical_layout=input_tensor.logical_layout,\n                output_logical_layout=model_ir.tensors[str(op.outputs[0])].logical_layout,\n                depthwise=True,\n            )\n            conv_ctor_lines = [\n                "torch.nn.Conv2d(",\n                f"    in_channels={depthwise_in_channels},",\n                f"    out_channels={int(weight_tensor.shape[0])},",\n                f"    kernel_size={_shape_literal(weight_tensor.shape[2:])},",\n                f"    stride={_shape_literal([int(options.get(\'strideH\', 1)), int(options.get(\'strideW\', 1))])},",\n                f"    padding={_shape_literal(conv_padding)},",\n                f"    dilation={_shape_literal([int(options.get(\'dilationHFactor\', 1)), int(options.get(\'dilationWFactor\', 1))])},",\n                f"    groups={depthwise_groups},",\n                f"    bias={str(bias_name != \'\')},",\n                ")",\n            ]\n            if use_conv_block:\n                module_init_lines.extend(\n                    [\n                        f"self.{attr_name} = _Conv2dBlock(",\n                        f"    in_channels={depthwise_in_channels},",\n                        f"    out_channels={int(weight_tensor.shape[0])},",\n                        f"    kernel_size={_shape_literal(weight_tensor.shape[2:])},",\n                        f"    stride={_shape_literal([int(options.get(\'strideH\', 1)), int(options.get(\'strideW\', 1))])},",\n                        f"    padding={_shape_literal(conv_padding)},",\n                        f"    dilation={_shape_literal([int(options.get(\'dilationHFactor\', 1)), int(options.get(\'dilationWFactor\', 1))])},",\n                        f"    groups={depthwise_groups},",\n                        f"    bias={str(bias_name != \'\')},",\n                        f"    pad={fused_block_pad if fused_block_pad is not None else (repr(conv_same_pad_arg) if conv_same_pad_arg is not None else \'None\')},",\n                        f"    activation={fused_block_activation!r},",\n                        f"    negative_slope={repr(fused_block_negative_slope if fused_block_negative_slope is not None else 0.2)},",\n                        f"    pad_mode={fused_block_pad_mode!r},",\n                        f"    pad_value={fused_block_pad_value if fused_block_pad_value is not None else \'0.0\'},",\n                        ")",\n                    ]\n                )\n            else:\n                module_init_lines.extend(\n                    [\n                        f"self.{attr_name} = {conv_ctor_lines[0]}",\n                        *conv_ctor_lines[1:],\n                    ]\n                )\n        elif op_type == "TRANSPOSE_CONV":\n            module_init_lines.extend(\n                [\n                    f"self.{attr_name} = torch.nn.ConvTranspose2d(",\n                    f"    in_channels={int(weight_tensor.shape[0])},",\n                    f"    out_channels={int(weight_tensor.shape[1])},",\n                    f"    kernel_size={_shape_literal(weight_tensor.shape[2:])},",\n                    f"    stride={_shape_literal([int(options.get(\'strideH\', 1)), int(options.get(\'strideW\', 1))])},",\n                    f"    padding={_shape_literal([int(((int(weight_tensor.shape[2]) - 1)) // 2) if str(options.get(\'padding\', \'SAME\')).upper() != \'VALID\' else 0, int(((int(weight_tensor.shape[3]) - 1)) // 2) if str(options.get(\'padding\', \'SAME\')).upper() != \'VALID\' else 0])},",\n                    f"    bias={str(bias_name != \'\')},",\n                    ")",\n                ]\n            )\n        elif op_type == "CONV_3D":\n            conv3d_in_channels, conv3d_groups, conv3d_out_channels, conv3d_kernel_size = _infer_conv3d_ctor_params(\n                input_shape=list(input_tensor.shape),\n                output_shape=list(model_ir.tensors[str(op.outputs[0])].shape),\n                weight_shape=list(weight_tensor.shape),\n                options=options,\n                input_logical_layout=normalize_logical_layout(input_tensor.logical_layout),\n                output_logical_layout=normalize_logical_layout(\n                    model_ir.tensors[str(op.outputs[0])].logical_layout\n                ),\n            )\n            module_init_lines.extend(\n                [\n                    f"self.{attr_name} = torch.nn.Conv3d(",\n                    f"    in_channels={conv3d_in_channels},",\n                    f"    out_channels={conv3d_out_channels},",\n                    f"    kernel_size={_shape_literal(conv3d_kernel_size)},",\n                    f"    stride={_shape_literal([int(options.get(\'strideD\', 1)), int(options.get(\'strideH\', 1)), int(options.get(\'strideW\', 1))])},",\n                    f"    padding={_shape_literal([int(((int(conv3d_kernel_size[0]) - 1) * int(options.get(\'dilationDFactor\', 1))) // 2) if str(options.get(\'padding\', \'SAME\')).upper() != \'VALID\' else 0, int(((int(conv3d_kernel_size[1]) - 1) * int(options.get(\'dilationHFactor\', 1))) // 2) if str(options.get(\'padding\', \'SAME\')).upper() != \'VALID\' else 0, int(((int(conv3d_kernel_size[2]) - 1) * int(options.get(\'dilationWFactor\', 1))) // 2) if str(options.get(\'padding\', \'SAME\')).upper() != \'VALID\' else 0])},",\n                    f"    dilation={_shape_literal([int(options.get(\'dilationDFactor\', 1)), int(options.get(\'dilationHFactor\', 1)), int(options.get(\'dilationWFactor\', 1))])},",\n                    f"    groups={conv3d_groups},",\n                    f"    bias={str(bias_name != \'\')},",\n                    ")",\n                ]\n            )\n        elif op_type == "CONV_3D_TRANSPOSE":\n            conv3dt_in_channels, conv3dt_out_channels, conv3dt_kernel_size, conv3dt_groups = _infer_conv3d_transpose_ctor_params(\n                input_shape=list(input_tensor.shape),\n                output_shape=list(model_ir.tensors[str(op.outputs[0])].shape),\n                weight_shape=list(weight_tensor.shape),\n                options=options,\n                input_logical_layout=normalize_logical_layout(input_tensor.logical_layout),\n                output_logical_layout=normalize_logical_layout(\n                    model_ir.tensors[str(op.outputs[0])].logical_layout\n                ),\n            )\n            module_init_lines.extend(\n                [\n                    f"self.{attr_name} = torch.nn.ConvTranspose3d(",\n                    f"    in_channels={conv3dt_in_channels},",\n                    f"    out_channels={conv3dt_out_channels},",\n                    f"    kernel_size={_shape_literal(conv3dt_kernel_size)},",\n                    f"    stride={_shape_literal([int(options.get(\'strideD\', 1)), int(options.get(\'strideH\', 1)), int(options.get(\'strideW\', 1))])},",\n                    f"    padding={_shape_literal([int(((int(conv3dt_kernel_size[0]) - 1)) // 2) if str(options.get(\'padding\', \'SAME\')).upper() != \'VALID\' else 0, int(((int(conv3dt_kernel_size[1]) - 1)) // 2) if str(options.get(\'padding\', \'SAME\')).upper() != \'VALID\' else 0, int(((int(conv3dt_kernel_size[2]) - 1)) // 2) if str(options.get(\'padding\', \'SAME\')).upper() != \'VALID\' else 0])},",\n                    f"    groups={conv3dt_groups},",\n                    f"    bias={str(bias_name != \'\')},",\n                    ")",\n                ]\n            )\n        elif op_type == "FULLY_CONNECTED":\n            module_init_lines.extend(\n                [\n                    f"self.{attr_name} = torch.nn.Linear(",\n                    f"    in_features={int(weight_tensor.shape[1])},",\n                    f"    out_features={int(weight_tensor.shape[0])},",\n                    f"    bias={str(bias_name != \'\')},",\n                    ")",\n                ]\n            )\n        elif op_type == "PRELU":\n            num_parameters = 1\n            if isinstance(weight_tensor.data, np.ndarray):\n                num_parameters = max(1, int(np.asarray(weight_tensor.data).size))\n            elif len(list(weight_tensor.shape)) > 0:\n                shape_values = [int(v) for v in list(weight_tensor.shape) if int(v) > 0]\n                if len(shape_values) > 0:\n                    num_parameters = max(1, int(np.prod(shape_values, dtype=np.int64)))\n            module_init_lines.extend(\n                [\n                    f"self.{attr_name} = torch.nn.PReLU(",\n                    f"    num_parameters={num_parameters},",\n                    "    init=0.0,",\n                    ")",\n                ]\n            )\n        load_specs.append((f"{attr_name}.conv.weight" if use_conv_block else f"{attr_name}.weight", str(weight_name)))\n        if bias_name != "":\n            load_specs.append((f"{attr_name}.conv.bias" if use_conv_block else f"{attr_name}.bias", str(bias_name)))\n\n    for op in model_ir.operators:\n        op_type = str(op.op_type)\n        if op_type not in {"PAD", "PADV2", "MIRROR_PAD"}:\n            continue\n        if len(op.inputs) >= 2:\n            paddings_tensor = model_ir.tensors.get(str(op.inputs[1]), None)\n            if paddings_tensor is not None and isinstance(paddings_tensor.data, np.ndarray):\n                inlined_constant_tensor_names.add(str(op.inputs[1]))\n        if op_type == "PADV2" and len(op.inputs) >= 3:\n            value_tensor = model_ir.tensors.get(str(op.inputs[2]), None)\n            if value_tensor is not None and isinstance(value_tensor.data, np.ndarray):\n                inlined_constant_tensor_names.add(str(op.inputs[2]))\n    for tensor_name, tensor in model_ir.tensors.items():\n        if str(tensor_name) in module_param_tensor_names:\n            continue\n        if _is_small_inline_constant_tensor(tensor):\n            inlined_constant_tensor_names.add(str(tensor_name))\n\n    buffer_attr_names = _build_buffer_attr_name_map(\n        model_ir=model_ir,\n        tensor_storage_name_map=tensor_storage_name_map,\n        excluded_tensor_names=module_param_tensor_names | submodule_state_tensor_names | inlined_constant_tensor_names,\n    )\n    buffer_attr_name_to_tensor_name = {\n        str(attr_name): str(tensor_name) for tensor_name, attr_name in buffer_attr_names.items()\n    }\n    buffer_init_lines: List[str] = []\n    for tensor_name, attr_name in buffer_attr_names.items():\n        tensor = model_ir.tensors[str(tensor_name)]\n        dtype_name = str(tensor.dtype).upper()\n        shape_values = [int(v) for v in list(tensor.shape)]\n        if bool(tensor.is_variable):\n            buffer_init_lines.append(\n                f"self.register_parameter({attr_name!r}, torch.nn.Parameter(torch.zeros({repr(shape_values)}, dtype={_torch_dtype_literal(dtype_name)}), requires_grad=False))"\n            )\n        else:\n            buffer_init_lines.append(\n                f"self.register_buffer({attr_name!r}, torch.zeros({repr(shape_values)}, dtype={_torch_dtype_literal(dtype_name)}), persistent=True)"\n            )\n        load_specs.append((str(attr_name), str(tensor_name)))\n\n    _binary_trailing_axis_constant_buffer_alias_shape = lambda tensor_name, other_tensor_name: _binary_trailing_axis_constant_buffer_alias_shape_for_codegen(\n        model_ir=model_ir,\n        producer_index=producer_index,\n        inlined_constant_tensor_names=inlined_constant_tensor_names,\n        tensor_name=tensor_name,\n        other_tensor_name=other_tensor_name,\n    )\n    _channel_first_rank4_constant_buffer_alias_shape = lambda tensor_name: _channel_first_rank4_constant_buffer_alias_shape_for_codegen(\n        model_ir=model_ir,\n        producer_index=producer_index,\n        inlined_constant_tensor_names=inlined_constant_tensor_names,\n        tensor_name=tensor_name,\n    )\n    _constant_permute_for_broadcast = lambda tensor_name, other_tensor_name: _constant_permute_for_broadcast_for_codegen(\n        model_ir=model_ir,\n        tensor_name=tensor_name,\n        other_tensor_name=other_tensor_name,\n    )\n\n    binary_constant_buffer_alias_attr_names: Dict[Tuple[str, Tuple[int, ...]], str] = {}\n    binary_constant_buffer_alias_exprs: Dict[Tuple[str, str], str] = {}\n    channel_first_constant_buffer_alias_attr_names: Dict[Tuple[str, Tuple[int, ...]], str] = {}\n    channel_first_constant_buffer_alias_exprs: Dict[str, str] = {}\n    channel_first_constant_buffer_alias_refresh_specs: List[Tuple[str, str, List[int]]] = []\n    permuted_constant_buffer_alias_attr_names: Dict[Tuple[str, Tuple[int, ...]], str] = {}\n    permuted_constant_buffer_alias_exprs: Dict[Tuple[str, Tuple[int, ...]], str] = {}\n    permuted_constant_buffer_alias_refresh_specs: List[Tuple[str, str, List[int]]] = []\n    transposed_constant_buffer_alias_attr_names: Dict[Tuple[str, Tuple[int, ...]], str] = {}\n    transposed_constant_buffer_alias_exprs: Dict[str, str] = {}\n    transposed_constant_buffer_alias_refresh_specs: List[Tuple[str, str]] = []\n    used_buffer_attr_names: Set[str] = set(buffer_attr_names.values())\n    for op in model_ir.operators:\n        if str(op.op_type) not in _DIRECT_CODEGEN_BINARY_FUNCTIONS:\n            continue\n        input_names = [str(name) for name in list(op.inputs)]\n        if len(input_names) < 2:\n            continue\n        lhs_name = str(input_names[0])\n        rhs_name = str(input_names[1])\n        for constant_name, other_name in ((lhs_name, rhs_name), (rhs_name, lhs_name)):\n            tensor = model_ir.tensors.get(str(constant_name), None)\n            if tensor is None:\n                continue\n            dtype_name = str(tensor.dtype).upper()\n            alias_shape = _binary_trailing_axis_constant_buffer_alias_shape(\n                constant_name,\n                other_name,\n            )\n            if alias_shape is not None:\n                alias_key = (str(constant_name), tuple(int(v) for v in alias_shape))\n                attr_name = binary_constant_buffer_alias_attr_names.get(alias_key, None)\n                if attr_name is None:\n                    storage_name = tensor_storage_name_map.get(str(constant_name), str(constant_name))\n                    shape_suffix = "x".join(str(int(v)) for v in alias_shape)\n                    base_attr_name = _sanitize_python_identifier(\n                        f"const_{storage_name}_broadcast_{shape_suffix}",\n                        prefix="const",\n                    )\n                    attr_name = _make_unique_identifier(base_attr_name, used_buffer_attr_names)\n                    binary_constant_buffer_alias_attr_names[alias_key] = str(attr_name)\n                    if bool(tensor.is_variable):\n                        buffer_init_lines.append(\n                            f"self.register_parameter({attr_name!r}, torch.nn.Parameter(torch.zeros({repr(alias_shape)}, dtype={_torch_dtype_literal(dtype_name)}), requires_grad=False))"\n                        )\n                    else:\n                        buffer_init_lines.append(\n                            f"self.register_buffer({attr_name!r}, torch.zeros({repr(alias_shape)}, dtype={_torch_dtype_literal(dtype_name)}), persistent=True)"\n                        )\n                    load_specs.append((str(attr_name), str(constant_name)))\n                binary_constant_buffer_alias_exprs[(str(constant_name), str(other_name))] = f"self.{attr_name}"\n            channel_first_alias_shape = _channel_first_rank4_constant_buffer_alias_shape(\n                constant_name,\n            )\n            if channel_first_alias_shape is None:\n                source_attr_name = buffer_attr_names.get(str(constant_name), None)\n            else:\n                source_attr_name = buffer_attr_names.get(str(constant_name), None)\n                if source_attr_name is not None:\n                    alias_key = (str(constant_name), tuple(int(v) for v in channel_first_alias_shape))\n                    channel_first_attr_name = channel_first_constant_buffer_alias_attr_names.get(alias_key, None)\n                    if channel_first_attr_name is None:\n                        storage_name = tensor_storage_name_map.get(str(constant_name), str(constant_name))\n                        shape_suffix = "x".join(str(int(v)) for v in channel_first_alias_shape)\n                        base_attr_name = _shorten_generated_python_identifier(\n                            f"const_{storage_name}_channel_first_{shape_suffix}",\n                            prefix="const",\n                        )\n                        channel_first_attr_name = _make_unique_identifier(base_attr_name, used_buffer_attr_names)\n                        channel_first_constant_buffer_alias_attr_names[alias_key] = str(channel_first_attr_name)\n                        buffer_init_lines.append(\n                            f"self.register_buffer({channel_first_attr_name!r}, torch.zeros({repr(channel_first_alias_shape)}, dtype={_torch_dtype_literal(dtype_name)}), persistent=False)"\n                        )\n                        channel_first_constant_buffer_alias_refresh_specs.append(\n                            (\n                                str(channel_first_attr_name),\n                                str(source_attr_name),\n                                [int(v) for v in channel_first_alias_shape],\n                            )\n                        )\n                    channel_first_constant_buffer_alias_exprs[str(constant_name)] = f"self.{channel_first_attr_name}"\n\n            source_attr_name = buffer_attr_names.get(str(constant_name), None)\n            permute_for_broadcast = _constant_permute_for_broadcast(constant_name, other_name)\n            if source_attr_name is not None and permute_for_broadcast is not None:\n                tensor_shape = [int(v) for v in list(tensor.shape)]\n                permuted_shape = [int(tensor_shape[int(idx)]) for idx in list(permute_for_broadcast)]\n                alias_key = (str(constant_name), tuple(int(v) for v in list(permute_for_broadcast)))\n                permuted_attr_name = permuted_constant_buffer_alias_attr_names.get(alias_key, None)\n                if permuted_attr_name is None:\n                    storage_name = tensor_storage_name_map.get(str(constant_name), str(constant_name))\n                    perm_suffix = "_".join(str(int(v)) for v in list(permute_for_broadcast))\n                    shape_suffix = "x".join(str(int(v)) for v in list(permuted_shape))\n                    base_attr_name = _shorten_generated_python_identifier(\n                        f"const_{storage_name}_perm_{perm_suffix}_{shape_suffix}",\n                        prefix="const",\n                    )\n                    permuted_attr_name = _make_unique_identifier(base_attr_name, used_buffer_attr_names)\n                    permuted_constant_buffer_alias_attr_names[alias_key] = str(permuted_attr_name)\n                    buffer_init_lines.append(\n                        f"self.register_buffer({permuted_attr_name!r}, torch.zeros({repr(permuted_shape)}, dtype={_torch_dtype_literal(dtype_name)}), persistent=False)"\n                    )\n                    permuted_constant_buffer_alias_refresh_specs.append(\n                        (\n                            str(permuted_attr_name),\n                            str(source_attr_name),\n                            [int(v) for v in list(permute_for_broadcast)],\n                        )\n                    )\n                permuted_constant_buffer_alias_exprs[(str(constant_name), tuple(int(v) for v in list(permute_for_broadcast)))] = f"self.{permuted_attr_name}"\n    for op in model_ir.operators:\n        if str(op.op_type) != "BATCH_MATMUL" or len(op.inputs) < 2:\n            continue\n        input_names = [str(name) for name in list(op.inputs[:2])]\n        transpose_flags = [\n            bool(op.options.get("adjX", False)),\n            bool(op.options.get("adjY", False)),\n        ]\n        for input_name, requires_transpose in zip(input_names, transpose_flags):\n            if not bool(requires_transpose):\n                continue\n            tensor = model_ir.tensors.get(str(input_name), None)\n            if (\n                tensor is None\n                or not isinstance(tensor.data, np.ndarray)\n                or bool(tensor.is_variable)\n                or str(input_name) in model_ir.inputs\n                or str(input_name) in producer_index\n                or str(input_name) in inlined_constant_tensor_names\n            ):\n                continue\n            tensor_shape = [int(v) for v in list(tensor.shape)]\n            if len(tensor_shape) < 2:\n                continue\n            source_attr_name = buffer_attr_names.get(str(input_name), None)\n            if source_attr_name is None:\n                continue\n            alias_shape = list(tensor_shape[:-2]) + [int(tensor_shape[-1]), int(tensor_shape[-2])]\n            alias_key = (str(input_name), tuple(int(v) for v in alias_shape))\n            attr_name = transposed_constant_buffer_alias_attr_names.get(alias_key, None)\n            if attr_name is None:\n                dtype_name = str(tensor.dtype).upper()\n                storage_name = tensor_storage_name_map.get(str(input_name), str(input_name))\n                shape_suffix = "x".join(str(int(v)) for v in alias_shape)\n                base_attr_name = _shorten_generated_python_identifier(\n                    f"const_{storage_name}_transpose_last_two_{shape_suffix}",\n                    prefix="const",\n                )\n                attr_name = _make_unique_identifier(base_attr_name, used_buffer_attr_names)\n                transposed_constant_buffer_alias_attr_names[alias_key] = str(attr_name)\n                buffer_init_lines.append(\n                    f"self.register_buffer({attr_name!r}, torch.zeros({repr(alias_shape)}, dtype={_torch_dtype_literal(dtype_name)}), persistent=False)"\n                )\n                transposed_constant_buffer_alias_refresh_specs.append(\n                    (str(attr_name), str(source_attr_name))\n                )\n            transposed_constant_buffer_alias_exprs[str(input_name)] = f"self.{attr_name}"\n\n    _tensor_expr = lambda tensor_name: _tensor_expr_for_codegen(\n        model_ir=model_ir,\n        producer_index=producer_index,\n        tensor_expr_aliases=tensor_expr_aliases,\n        channel_first_tensor_expr_aliases=channel_first_tensor_expr_aliases,\n        buffer_attr_names=buffer_attr_names,\n        runtime_imports=runtime_imports,\n        tensor_var_names=tensor_var_names,\n        tensor_name=tensor_name,\n    )\n    _derived_local_var_name = lambda base_name, prefix="t": _derived_local_var_name_for_codegen(\n        synthetic_local_var_names=synthetic_local_var_names,\n        used_local_var_names=used_local_var_names,\n        base_name=base_name,\n        prefix=prefix,\n    )\n\n    _tensor_expr_for_channel_first_bridge = lambda tensor_name, perm: _tensor_expr_for_channel_first_bridge_for_codegen(\n        model_ir=model_ir,\n        channel_first_tensor_expr_aliases=channel_first_tensor_expr_aliases,\n        tensor_name=tensor_name,\n        perm=perm,\n    )\n    _channel_first_reduction_plan = lambda op, input_name: _channel_first_reduction_plan_for_codegen(\n        model_ir=model_ir,\n        channel_first_tensor_expr_aliases=channel_first_tensor_expr_aliases,\n        op=op,\n        input_name=input_name,\n    )\n    _normalized_constant_reduction_axes = lambda axis_values, rank: _normalized_constant_reduction_axes_for_codegen(\n        axis_values=axis_values,\n        rank=rank,\n    )\n    _direct_mean_reduction_expr = lambda *, input_expr, axes, input_rank, keepdims: _direct_mean_reduction_expr_for_codegen(\n        normalized_constant_reduction_axes_fn=_normalized_constant_reduction_axes,\n        input_expr=input_expr,\n        axes=axes,\n        input_rank=input_rank,\n        keepdims=keepdims,\n    )\n\n    _channel_first_passthrough_input_expr = lambda tensor_name: _channel_first_passthrough_input_expr_for_codegen(\n        model_ir=model_ir,\n        producer_index=producer_index,\n        channel_first_tensor_expr_aliases=channel_first_tensor_expr_aliases,\n        tensor_expr_fn=_tensor_expr,\n        tensor_name=tensor_name,\n    )\n    _can_resolve_channel_first_expr_statically = lambda tensor_name, seen_names=None: _can_resolve_channel_first_expr_statically_for_codegen(\n        model_ir=model_ir,\n        producer_index=producer_index,\n        channel_first_tensor_expr_aliases=channel_first_tensor_expr_aliases,\n        direct_codegen_unary_expressions=_DIRECT_CODEGEN_UNARY_EXPRESSIONS,\n        tensor_name=tensor_name,\n        seen_names=seen_names,\n    )\n    _can_emit_channel_first_shape_preserving_unary_op = lambda op: _can_emit_channel_first_shape_preserving_unary_op_for_codegen(\n        model_ir=model_ir,\n        direct_codegen_unary_expressions=_DIRECT_CODEGEN_UNARY_EXPRESSIONS,\n        tensor_shape_list_fn=_tensor_shape_list,\n        can_resolve_channel_first_expr_statically_fn=lambda tensor_name: _can_resolve_channel_first_expr_statically(tensor_name),\n        op=op,\n    )\n\n    _all_consumers_are_channel_first_binary_ops = lambda output_name: _all_consumers_are_channel_first_binary_ops_for_codegen(\n        model_ir=model_ir,\n        consumer_index=consumer_index,\n        direct_codegen_binary_functions=_DIRECT_CODEGEN_BINARY_FUNCTIONS,\n        can_emit_channel_first_binary_op_fn=lambda op: _can_emit_channel_first_binary_op(op),\n        output_name=output_name,\n    )\n    _can_omit_materialized_channel_last_alias_recursive = lambda output_name, seen_names: _can_omit_materialized_channel_last_alias_recursive_for_codegen(\n        model_ir=model_ir,\n        consumer_index=consumer_index,\n        direct_codegen_unary_expressions=_DIRECT_CODEGEN_UNARY_EXPRESSIONS,\n        tensor_shape_list_fn=_tensor_shape_list,\n        channel_first_reduction_plan_fn=_channel_first_reduction_plan,\n        can_emit_channel_first_shape_preserving_unary_op_fn=_can_emit_channel_first_shape_preserving_unary_op,\n        can_emit_channel_first_binary_op_fn=lambda op: _can_emit_channel_first_binary_op(op),\n        can_resolve_channel_first_expr_statically_fn=lambda tensor_name: _can_resolve_channel_first_expr_statically(tensor_name),\n        conv2d_input_pre_permute_fn=_conv2d_input_pre_permute,\n        output_name=output_name,\n        seen_names=seen_names,\n    )\n    _can_omit_materialized_channel_last_alias = lambda output_name: _can_omit_materialized_channel_last_alias_recursive(str(output_name), set())\n\n    _channel_first_binary_input_expr = lambda tensor_name, other_tensor_name: _channel_first_binary_input_expr_for_codegen(\n        model_ir=model_ir,\n        channel_first_tensor_expr_aliases=channel_first_tensor_expr_aliases,\n        channel_first_constant_buffer_alias_exprs=channel_first_constant_buffer_alias_exprs,\n        permuted_constant_buffer_alias_exprs=permuted_constant_buffer_alias_exprs,\n        scalar_literal_expr_fn=_scalar_literal_expr,\n        tensor_shape_list_fn=_tensor_shape_list,\n        tensor_expr_fn=_tensor_expr,\n        channel_first_passthrough_input_expr_fn=_channel_first_passthrough_input_expr,\n        tensor_name=tensor_name,\n        other_tensor_name=other_tensor_name,\n    )\n\n    _channel_first_constant_expr_for_buffer_attr = lambda buffer_expr, target_shape: _channel_first_constant_expr_for_buffer_attr_for_codegen(\n        buffer_attr_name_to_tensor_name=buffer_attr_name_to_tensor_name,\n        channel_first_constant_buffer_alias_exprs=channel_first_constant_buffer_alias_exprs,\n        channel_first_rank4_constant_buffer_alias_shape_fn=_channel_first_rank4_constant_buffer_alias_shape,\n        buffer_expr=buffer_expr,\n        target_shape=target_shape,\n    )\n    _permuted_constant_expr_for_tensor_name = lambda tensor_name, perm: _permuted_constant_expr_for_tensor_name_for_codegen(\n        permuted_constant_buffer_alias_exprs=permuted_constant_buffer_alias_exprs,\n        tensor_name=tensor_name,\n        perm=perm,\n    )\n    _transposed_constant_expr_for_tensor_name = lambda tensor_name: _transposed_constant_expr_for_tensor_name_for_codegen(\n        transposed_constant_buffer_alias_exprs=transposed_constant_buffer_alias_exprs,\n        tensor_name=tensor_name,\n    )\n\n    _can_emit_channel_first_binary_op = lambda op: _can_emit_channel_first_binary_op_for_codegen(\n        model_ir=model_ir,\n        tensor_shape_list_fn=_tensor_shape_list,\n        channel_first_shape_for_tensor_fn=_channel_first_shape_for_tensor,\n        scalar_literal_expr_fn=_scalar_literal_expr,\n        can_omit_materialized_channel_last_alias_fn=_can_omit_materialized_channel_last_alias,\n        channel_first_binary_input_expr_fn=_channel_first_binary_input_expr,\n        op=op,\n    )\n\n    producer_by_output_name: Dict[str, OperatorIR] = {}\n    for operator in model_ir.operators:\n        for output_name in operator.outputs:\n            producer_by_output_name[str(output_name)] = operator\n\n    _shape_tensor_length = lambda tensor_name: _shape_tensor_length_for_codegen(\n        model_ir=model_ir,\n        tensor_name=tensor_name,\n    )\n    _reconstruct_shape_scalar_expr = lambda tensor_name, seen=None: _reconstruct_shape_scalar_expr_for_codegen(\n        model_ir=model_ir,\n        producer_by_output_name=producer_by_output_name,\n        tensor_exact_static_shape_list_fn=_tensor_exact_static_shape_list,\n        tensor_expr_fn=_tensor_expr,\n        runtime_imports=runtime_imports,\n        tensor_name=tensor_name,\n        seen=seen,\n    )\n    _reconstruct_shape_list_expr = lambda tensor_name, seen=None: _reconstruct_shape_list_expr_for_codegen(\n        model_ir=model_ir,\n        producer_by_output_name=producer_by_output_name,\n        tensor_exact_static_shape_list_fn=_tensor_exact_static_shape_list,\n        tensor_expr_fn=_tensor_expr,\n        runtime_imports=runtime_imports,\n        tensor_name=tensor_name,\n        seen=seen,\n    )\n\n    _match_if_axis0_tensor_mux_slice = lambda op: _match_if_axis0_tensor_mux_slice_for_codegen(\n        model_ir=model_ir,\n        producer_by_output_name=producer_by_output_name,\n        op=op,\n    )\n    _binary_operand_expr = lambda tensor_name, other_tensor_name: _binary_operand_expr_for_codegen(\n        model_ir=model_ir,\n        binary_constant_buffer_alias_exprs=binary_constant_buffer_alias_exprs,\n        channel_first_constant_buffer_alias_exprs=channel_first_constant_buffer_alias_exprs,\n        channel_first_tensor_expr_aliases=channel_first_tensor_expr_aliases,\n        tensor_expr_fn=_tensor_expr,\n        channel_first_rank4_constant_buffer_alias_shape_fn=_channel_first_rank4_constant_buffer_alias_shape,\n        channel_first_shape_for_tensor_fn=_channel_first_shape_for_tensor,\n        constant_permute_for_broadcast_fn=_constant_permute_for_broadcast,\n        permuted_constant_expr_for_tensor_name_fn=_permuted_constant_expr_for_tensor_name,\n        tensor_name=tensor_name,\n        other_tensor_name=other_tensor_name,\n    )\n\n    runtime_shape_uncertain_tensors: Set[str] = set()\n\n    _binary_runtime_shape_passthrough_operand = lambda lhs_name, rhs_name: _binary_runtime_shape_passthrough_operand_for_codegen(\n        model_ir=model_ir,\n        runtime_shape_uncertain_tensors=runtime_shape_uncertain_tensors,\n        lhs_name=lhs_name,\n        rhs_name=rhs_name,\n    )\n    _binary_requires_runtime_alignment = lambda lhs_name, rhs_name, output_name: _binary_requires_runtime_alignment_for_codegen(\n        model_ir=model_ir,\n        runtime_shape_uncertain_tensors=runtime_shape_uncertain_tensors,\n        lhs_name=lhs_name,\n        rhs_name=rhs_name,\n        output_name=output_name,\n    )\n\n    if _is_sequential_single_input_graph():\n        tensor_var_names[str(model_ir.inputs[0])] = "x"\n        for op in model_ir.operators:\n            if str(op.op_type) == "SHAPE":\n                continue\n            tensor_var_names[str(op.outputs[0])] = "x"\n    _has_channel_last_consumer_hint_for_same_shape_transpose = lambda op: _has_channel_last_consumer_hint_for_same_shape_transpose_for_codegen(\n        model_ir=model_ir,\n        consumer_index=consumer_index,\n        op=op,\n    )\n    _is_batchless_rank3_public_output_transpose = lambda op: _is_batchless_rank3_public_output_transpose_for_codegen(\n        model_ir=model_ir,\n        producer_index=producer_index,\n        batchless_rank3_public_boundary_names=batchless_rank3_public_boundary_names,\n        op=op,\n    )\n\n    runtime_imports = context.runtime_imports\n    runtime_imports.update({\n        "resolve_named_input_value",\n        "load_generated_weights",\n    })\n    batchless_rank3_public_boundary_names = {\n        str(name)\n        for name in list(model_ir.metadata.get("batchless_rank3_public_boundary_names", []))\n    }\n    forward_lines = context.forward_lines\n    for op_index, op in enumerate(model_ir.operators):\n        if int(op_index) in skipped_op_indices:\n            continue\n        op_type = str(op.op_type)\n        outputs = [str(v) for v in op.outputs]\n        output_vars = [tensor_var_names[str(name)] for name in outputs]\n        output_target_shape = _target_shape_literal(outputs[0]) if len(outputs) == 1 else "None"\n        if (\n            op_type in {"GATHER", "GATHER_ND", "RESHAPE", "SLICE", "STRIDED_SLICE", "CONCATENATION"}\n            and any(str(input_name) in runtime_shape_uncertain_tensors for input_name in op.inputs)\n        ):\n            runtime_shape_uncertain_tensors.update(outputs)\n        affine_layer_norm_spec = affine_layer_norm_specs.get(int(op_index), None)\n        if affine_layer_norm_spec is not None:\n            forward_lines.append(\n                f"{output_vars[0]} = self.{affine_layer_norm_spec[\'attr_name\']}({_tensor_expr(str(affine_layer_norm_spec[\'input_name\']))})"\n            )\n            continue\n        if _emit_native_direct_module_op_for_codegen(\n            model_ir=model_ir,\n            op=op,\n            op_index=int(op_index),\n            outputs=outputs,\n            output_vars=output_vars,\n            output_target_shape=output_target_shape,\n            op_module_attr_names=op_module_attr_names,\n            fused_module_specs=fused_module_specs,\n            conv_module_pad_specs=conv_module_pad_specs,\n            tensor_var_names=tensor_var_names,\n            channel_first_tensor_expr_aliases=channel_first_tensor_expr_aliases,\n            runtime_imports=runtime_imports,\n            forward_lines=forward_lines,\n            tensor_expr_fn=_tensor_expr,\n            tensor_expr_for_channel_first_bridge_fn=_tensor_expr_for_channel_first_bridge,\n            all_consumers_are_channel_first_binary_ops_fn=_all_consumers_are_channel_first_binary_ops,\n            can_omit_materialized_channel_last_alias_fn=_can_omit_materialized_channel_last_alias,\n            derived_local_var_name_fn=_derived_local_var_name,\n            emit_module_output_expr_fn=_emit_module_output_expr,\n            target_shape_literal_fn=_target_shape_literal,\n            conv2d_input_pre_permute_fn=_conv2d_input_pre_permute,\n            can_emit_direct_module_call_fn=_can_emit_direct_module_call,\n            activation_lines_fn=_activation_lines,\n            emit_maybe_aligned_expr_fn=_emit_maybe_aligned_expr,\n            tensor_shape_list_fn=_tensor_shape_list,\n            should_skip_align_for_shape_preserving_unary_fn=_should_skip_align_for_shape_preserving_unary,\n        ):\n            continue\n        if _emit_native_binary_op_for_codegen(\n            model_ir=model_ir,\n            op=op,\n            op_index=int(op_index),\n            outputs=outputs,\n            output_vars=output_vars,\n            output_target_shape=output_target_shape,\n            channel_first_tensor_expr_aliases=channel_first_tensor_expr_aliases,\n            runtime_imports=runtime_imports,\n            forward_lines=forward_lines,\n            runtime_shape_uncertain_tensors=runtime_shape_uncertain_tensors,\n            tensor_dtype_name_fn=_tensor_dtype_name,\n            binary_operand_expr_fn=_binary_operand_expr,\n            scalar_literal_expr_fn=_scalar_literal_expr,\n            can_emit_channel_first_binary_op_fn=_can_emit_channel_first_binary_op,\n            channel_first_binary_input_expr_fn=_channel_first_binary_input_expr,\n            derived_local_var_name_fn=_derived_local_var_name,\n            can_omit_materialized_channel_last_alias_fn=_can_omit_materialized_channel_last_alias,\n            target_shape_literal_fn=_target_shape_literal,\n            emit_maybe_aligned_expr_fn=_emit_maybe_aligned_expr,\n            binary_runtime_shape_passthrough_operand_fn=_binary_runtime_shape_passthrough_operand,\n            binary_requires_runtime_alignment_fn=_binary_requires_runtime_alignment,\n            preferred_binary_alignment_anchor_fn=lambda lhs_name, rhs_name, output_name: _preferred_binary_alignment_anchor_for_codegen(\n                model_ir=model_ir,\n                lhs_name=lhs_name,\n                rhs_name=rhs_name,\n                output_name=output_name,\n            ),\n            activation_lines_fn=_activation_lines,\n        ):\n            continue\n        if _emit_native_unary_op_for_codegen(\n            model_ir=model_ir,\n            op=op,\n            outputs=outputs,\n            output_vars=output_vars,\n            channel_first_tensor_expr_aliases=channel_first_tensor_expr_aliases,\n            runtime_imports=runtime_imports,\n            forward_lines=forward_lines,\n            tensor_expr_fn=_tensor_expr,\n            channel_first_passthrough_input_expr_fn=_channel_first_passthrough_input_expr,\n            can_emit_channel_first_shape_preserving_unary_op_fn=_can_emit_channel_first_shape_preserving_unary_op,\n            derived_local_var_name_fn=_derived_local_var_name,\n            can_omit_materialized_channel_last_alias_fn=_can_omit_materialized_channel_last_alias,\n            target_shape_literal_fn=_target_shape_literal,\n            tensor_shape_list_fn=_tensor_shape_list,\n            should_skip_align_for_shape_preserving_unary_fn=_should_skip_align_for_shape_preserving_unary,\n            emit_maybe_aligned_expr_fn=_emit_maybe_aligned_expr,\n        ):\n            continue\n        if op_type == "GATHER":\n            axis = int(op.options.get("axis", 0))\n            batch_dims = int(op.options.get("batchDims", 0))\n            params_expr = _tensor_expr(str(op.inputs[0]))\n            gather_output_tensor = model_ir.tensors.get(str(outputs[0]), None)\n            gather_single_consumer_indices = consumer_index.get(str(outputs[0]), [])\n            if (\n                gather_output_tensor is not None\n                and len(gather_single_consumer_indices) == 1\n            ):\n                gather_consumer = model_ir.operators[int(gather_single_consumer_indices[0])]\n                gather_consumer_output_tensor = (\n                    model_ir.tensors.get(str(gather_consumer.outputs[0]), None)\n                    if len(gather_consumer.outputs) == 1\n                    else None\n                )\n                if (\n                    str(gather_consumer.op_type) == "RESHAPE"\n                    and gather_consumer_output_tensor is not None\n                    and _is_suffix_flatten_gather_reshape(\n                        gather_output_tensor.shape,\n                        gather_consumer_output_tensor.shape,\n                    )\n                ):\n                    continue\n            if _should_elide_crd_to_dcr_gather_for_depth_to_space(\n                model_ir=model_ir,\n                params_name=str(op.inputs[0]),\n                indices_name=str(op.inputs[1]),\n                output_name=str(outputs[0]),\n                axis=axis,\n                batch_dims=batch_dims,\n            ):\n                inferred_shape = _tensor_shape_list(str(op.inputs[0]))\n                if _should_skip_align_for_shape_preserving_unary(str(op.inputs[0]), outputs[0]):\n                    forward_lines.append(f"{output_vars[0]} = {params_expr}")\n                else:\n                    forward_lines.append(\n                        f"{output_vars[0]} = {_emit_maybe_aligned_expr(output_name=outputs[0], expr=params_expr, inferred_shape=inferred_shape)}"\n                    )\n                continue\n            gather_pre_perm = _gather_input_pre_permute(\n                params_name=str(op.inputs[0]),\n                output_name=str(outputs[0]),\n                axis=axis,\n                batch_dims=batch_dims,\n            )\n            gather_input_rank = len(model_ir.tensors[str(op.inputs[0])].shape)\n            if gather_pre_perm is not None:\n                params_expr = f"{params_expr}.permute({\', \'.join(str(int(v)) for v in gather_pre_perm)}).contiguous()"\n                gather_input_rank = int(len(gather_pre_perm))\n            indices_expr = _tensor_expr(str(op.inputs[1]))\n            indices_name = str(op.inputs[1])\n            if _is_identity_nms_postprocess_gather(str(op.inputs[0]), indices_name):\n                inferred_shape = _tensor_shape_list(str(op.inputs[0]))\n                forward_lines.append(\n                    f"{output_vars[0]} = {_emit_maybe_aligned_expr(output_name=outputs[0], expr=params_expr, inferred_shape=inferred_shape)}"\n                )\n                continue\n            direct_indices_values = _constant_int_list(model_ir.tensors.get(indices_name, None))\n            direct_gather_expr = None\n            indices_tensor = model_ir.tensors.get(indices_name, None)\n            if direct_indices_values is not None:\n                direct_gather_expr = _direct_gather_expr(\n                    params_expr=params_expr,\n                    indices_values=direct_indices_values,\n                    indices_shape=indices_tensor.shape if indices_tensor is not None else None,\n                    axis=axis,\n                    batch_dims=batch_dims,\n                    input_rank=gather_input_rank,\n                )\n            if direct_gather_expr is None:\n                direct_gather_expr = _direct_dynamic_gather_expr(\n                    params_expr=params_expr,\n                    indices_expr=indices_expr,\n                    axis=axis,\n                    batch_dims=batch_dims,\n                    input_rank=gather_input_rank,\n                    indices_name=indices_name,\n                    indices_shape=(\n                        indices_tensor.shape\n                        if indices_tensor is not None\n                        else None\n                    ),\n                    indices_shape_signature=(\n                        indices_tensor.shape_signature\n                        if indices_tensor is not None\n                        else None\n                    ),\n                )\n            if direct_gather_expr is not None:\n                if "_reshape_gather_output(" in direct_gather_expr:\n                    runtime_imports.add("_reshape_gather_output")\n                if "_shape_tensor(" in direct_gather_expr:\n                    runtime_imports.add("_shape_tensor")\n                forward_lines.append(f"{output_vars[0]} = {direct_gather_expr}")\n            else:\n                runtime_imports.add("_apply_gather")\n                forward_lines.append(\n                    f"{output_vars[0]} = _apply_gather({params_expr}, {indices_expr}, axis={axis}, batch_dims={batch_dims}, target_shape={output_target_shape}, indices_name={indices_name!r})"\n                )\n            continue\n        if op_type == "GATHER_ND":\n            params_name = str(op.inputs[0])\n            params_expr = _tensor_expr(params_name)\n            indices_expr = _tensor_expr(str(op.inputs[1]))\n            params_tensor = model_ir.tensors.get(params_name, None)\n            public_layout_map = cast(\n                Dict[str, str],\n                model_ir.metadata.get("onnx_public_layout_map", {}),\n            )\n            if params_tensor is not None and params_name in model_ir.inputs:\n                params_layout = normalize_logical_layout(params_tensor.logical_layout)\n                public_layout = normalize_logical_layout(\n                    public_layout_map.get(params_name, LOGICAL_LAYOUT_UNKNOWN)\n                )\n                if (\n                    params_layout != LOGICAL_LAYOUT_UNKNOWN\n                    and public_layout != LOGICAL_LAYOUT_UNKNOWN\n                    and params_layout != public_layout\n                ):\n                    boundary_perm = logical_layout_permutation(\n                        source_layout=public_layout,\n                        target_layout=params_layout,\n                    )\n                    if boundary_perm is not None:\n                        runtime_imports.add("_torch_permute")\n                        params_expr = f"_torch_permute({params_expr}, {repr([int(v) for v in list(boundary_perm)])})"\n            if params_tensor is not None:\n                params_layout = normalize_logical_layout(params_tensor.logical_layout)\n                params_actual_shape = [int(v) for v in list(params_tensor.shape)]\n                params_semantic_shape = _tensor_shape_list(params_name)\n                params_rank = len(params_actual_shape)\n                if (\n                    params_semantic_shape is not None\n                    and params_rank in {3, 4, 5}\n                    and is_channel_first_logical_layout(params_layout)\n                ):\n                    cf_to_cl_perm = logical_layout_permutation(\n                        source_layout=params_layout,\n                        target_layout=channel_last_logical_layout(params_rank),\n                    )\n                    if (\n                        cf_to_cl_perm is not None\n                        and _shape_lists_equal(\n                            _permute_shape(params_actual_shape, cf_to_cl_perm),\n                            params_semantic_shape,\n                        )\n                    ):\n                        runtime_imports.add("_torch_permute")\n                        params_expr = f"_torch_permute({params_expr}, {repr([int(v) for v in list(cf_to_cl_perm)])})"\n                output_tensor_for_gather_nd = model_ir.tensors.get(str(outputs[0]), None)\n                output_shape_for_gather_nd = (\n                    [int(v) for v in list(output_tensor_for_gather_nd.shape)]\n                    if output_tensor_for_gather_nd is not None\n                    else _tensor_shape_list(str(outputs[0]))\n                )\n                if (\n                    output_shape_for_gather_nd is not None\n                    and params_rank in {3, 4, 5}\n                    and is_channel_first_logical_layout(params_layout)\n                ):\n                    cf_to_cl_perm = logical_layout_permutation(\n                        source_layout=params_layout,\n                        target_layout=channel_last_logical_layout(params_rank),\n                    )\n                    if cf_to_cl_perm is not None:\n                        actual_gather_nd_shape = _infer_gather_nd_shape(\n                            params_shape=params_actual_shape,\n                            indices_tensor_name=str(op.inputs[1]),\n                        )\n                        permuted_actual_shape = _permute_shape(params_actual_shape, cf_to_cl_perm)\n                        permuted_gather_nd_shape = _infer_gather_nd_shape(\n                            params_shape=permuted_actual_shape,\n                            indices_tensor_name=str(op.inputs[1]),\n                        )\n                        if (\n                            not _shape_lists_equal(actual_gather_nd_shape, output_shape_for_gather_nd)\n                            and _shape_lists_equal(permuted_gather_nd_shape, output_shape_for_gather_nd)\n                        ):\n                            runtime_imports.add("_torch_permute")\n                            params_expr = f"_torch_permute({params_expr}, {repr([int(v) for v in list(cf_to_cl_perm)])})"\n            gather_elements_axis_hint = None\n            gather_elements_indices_name = None\n            if str(op.inputs[1]).endswith("_gather_elements_coords"):\n                producer_op_idx = producer_index.get(str(op.inputs[1]), None)\n                if producer_op_idx is not None:\n                    producer_op = model_ir.operators[int(producer_op_idx)]\n                    if str(producer_op.op_type) == "CONCATENATION":\n                        gather_elements_axis_hint = next(\n                            (\n                                idx for idx, name in enumerate(producer_op.inputs)\n                                if str(name).endswith("_gather_elements_axis_coord")\n                            ),\n                            None,\n                        )\n                candidate_indices_name = str(op.inputs[1]).replace("_gather_elements_coords", "_gather_elements_indices_i32")\n                if candidate_indices_name in model_ir.tensors:\n                    gather_elements_indices_name = candidate_indices_name\n            if gather_elements_axis_hint is not None and gather_elements_indices_name is not None:\n                params_tensor = model_ir.tensors.get(str(op.inputs[0]), None)\n                params_rank = 0 if params_tensor is None else len(list(params_tensor.shape))\n                gather_axis = int(gather_elements_axis_hint)\n                gather_axis_cl = None\n                params_layout = (\n                    LOGICAL_LAYOUT_UNKNOWN\n                    if params_tensor is None\n                    else normalize_logical_layout(params_tensor.logical_layout)\n                )\n                if params_rank in {3, 4, 5} and is_channel_first_logical_layout(params_layout):\n                    perm_cf_to_cl = logical_layout_permutation(\n                        source_layout=channel_first_logical_layout(params_rank),\n                        target_layout=channel_last_logical_layout(params_rank),\n                    )\n                    if perm_cf_to_cl is not None and gather_axis in perm_cf_to_cl:\n                        gather_axis_cl = int(perm_cf_to_cl.index(gather_axis))\n                runtime_imports.add("_apply_gather_elements")\n                runtime_shape_uncertain_tensors.add(outputs[0])\n                forward_lines.append(\n                    f"{output_vars[0]} = _apply_gather_elements({params_expr}, {_tensor_expr(gather_elements_indices_name)}, axis_hint={gather_axis}, axis_hint_cl={repr(gather_axis_cl)}, target_shape={output_target_shape})"\n                )\n                continue\n            params_shape = (\n                [int(v) for v in list(params_tensor.shape)]\n                if params_tensor is not None and params_name in model_ir.inputs\n                else _tensor_shape_list(str(op.inputs[0]))\n            )\n            output_shape = _tensor_shape_list(outputs[0])\n            if (\n                params_tensor is not None\n                and params_shape is not None\n                and output_shape is not None\n                and len(list(params_shape)) in {3, 4, 5}\n            ):\n                params_layout = normalize_logical_layout(params_tensor.logical_layout)\n                expected_shape_from_params = _infer_gather_nd_shape(params_shape, str(op.inputs[1]))\n                if not _shape_lists_equal(expected_shape_from_params, output_shape):\n                    alternate_layout = None\n                    rank = len(list(params_shape))\n                    if is_channel_first_logical_layout(params_layout):\n                        alternate_layout = channel_last_logical_layout(rank)\n                    elif is_channel_last_logical_layout(params_layout):\n                        alternate_layout = channel_first_logical_layout(rank)\n                    if alternate_layout is not None:\n                        perm = logical_layout_permutation(\n                            source_layout=params_layout,\n                            target_layout=alternate_layout,\n                        )\n                        if perm is not None:\n                            resolved_perm = perm\n                            permuted_params_shape = _permute_shape(params_shape, resolved_perm)\n                            if (\n                                permuted_params_shape is not None\n                                and _shape_lists_equal(\n                                    _infer_gather_nd_shape(permuted_params_shape, str(op.inputs[1])),\n                                    output_shape,\n                                )\n                            ):\n                                params_expr = f"{params_expr}.permute({\', \'.join(str(int(v)) for v in resolved_perm)}).contiguous()"\n            runtime_imports.add("_apply_gather_nd")\n            forward_lines.append(\n                f"{output_vars[0]} = _apply_gather_nd({params_expr}, {indices_expr}, target_shape={output_target_shape})"\n            )\n            continue\n        if op_type == "CAST":\n            out_dtype = str(op.options.get("outDataType", "FLOAT32"))\n            forward_lines.append(\n                f"{output_vars[0]} = {_tensor_expr(str(op.inputs[0]))}.to(dtype={_torch_dtype_literal(out_dtype)})"\n            )\n            continue\n        if op_type == "LOCAL_RESPONSE_NORMALIZATION":\n            input_expr = _tensor_expr(str(op.inputs[0]))\n            radius = int(op.options.get("radius", 0))\n            size = max(1, int(radius) * 2 + 1)\n            alpha = float(op.options.get("alpha", 1.0))\n            beta = float(op.options.get("beta", 0.5))\n            bias = float(op.options.get("bias", 1.0))\n            inferred_shape = _tensor_shape_list(str(op.inputs[0]))\n            forward_lines.append(\n                f"{output_vars[0]} = {_emit_maybe_aligned_expr(output_name=outputs[0], expr=f\'F.local_response_norm({input_expr}, size={size}, alpha={alpha}, beta={beta}, k={bias})\', inferred_shape=inferred_shape)}"\n            )\n            continue\n        if op_type == "RESHAPE":\n            runtime_imports.add("_shape_list")\n            reshape_input_expr = _tensor_expr(str(op.inputs[0]))\n            reshape_input_tensor = model_ir.tensors.get(str(op.inputs[0]), None)\n            reshape_output_tensor = model_ir.tensors.get(str(outputs[0]), None)\n            reshape_is_lowered_onnx_flatten = "onnxFlattenAxis" in op.options\n            reshape_allow_zero = bool(op.options.get("allowZero", False))\n            reshape_input_shape = None if reshape_input_tensor is None else [int(v) for v in list(reshape_input_tensor.shape)]\n            reshape_output_shape = None if reshape_output_tensor is None else [int(v) for v in list(reshape_output_tensor.shape)]\n            reshape_output_preferred_shape = _preferred_reshape_target_values(reshape_output_tensor)\n            reshape_output_semantic_shape = (\n                [int(v) for v in list(reshape_output_preferred_shape)]\n                if reshape_output_preferred_shape is not None\n                else reshape_output_shape\n            )\n            reshape_input_layout = None if reshape_input_tensor is None else str(reshape_input_tensor.logical_layout)\n            reshape_output_layout = None if reshape_output_tensor is None else str(reshape_output_tensor.logical_layout)\n            if (\n                not reshape_is_lowered_onnx_flatten\n                and reshape_input_shape is not None\n                and len(reshape_input_shape) == 4\n                and reshape_output_semantic_shape is not None\n                and len(reshape_output_semantic_shape) in {2, 3}\n            ):\n                effective_layout = _infer_effective_rank4_runtime_layout(str(op.inputs[0]))\n                if effective_layout is not None:\n                    reshape_input_layout = effective_layout\n            reshape_special_plan = None\n            reshape_pre_perm = None\n            reshape_feature_last_target = None\n            reshape_channel_first_alias_shape = None\n            gather_flatten_direct_spec: Optional[Dict[str, Any]] = None\n            reshape_plain_singleton_axis_drop = _reshape_is_plain_singleton_axis_drop(\n                reshape_input_shape,\n                reshape_output_semantic_shape,\n            )\n            reshape_input_producer_idx = producer_index.get(str(op.inputs[0]), None)\n            if reshape_input_producer_idx is not None:\n                reshape_input_producer = model_ir.operators[int(reshape_input_producer_idx)]\n                if (\n                    str(reshape_input_producer.op_type) == "GATHER"\n                    and len(consumer_index.get(str(op.inputs[0]), [])) == 1\n                    and _is_suffix_flatten_gather_reshape(\n                        reshape_input_shape,\n                        reshape_output_semantic_shape,\n                    )\n                ):\n                    gather_axis = int(reshape_input_producer.options.get("axis", 0))\n                    gather_batch_dims = int(reshape_input_producer.options.get("batchDims", 0))\n                    gather_params_name = str(reshape_input_producer.inputs[0])\n                    gather_indices_name = str(reshape_input_producer.inputs[1])\n                    gather_params_expr = _tensor_expr(gather_params_name)\n                    gather_input_rank = len(model_ir.tensors[gather_params_name].shape)\n                    gather_pre_perm = _gather_input_pre_permute(\n                        params_name=gather_params_name,\n                        output_name=str(op.inputs[0]),\n                        axis=gather_axis,\n                        batch_dims=gather_batch_dims,\n                    )\n                    if gather_pre_perm is not None:\n                        gather_params_expr = (\n                            f"{gather_params_expr}.permute("\n                            f"{\', \'.join(str(int(v)) for v in gather_pre_perm)}).contiguous()"\n                        )\n                        gather_input_rank = int(len(gather_pre_perm))\n                    gather_flatten_direct_spec = {\n                        "params_expr": gather_params_expr,\n                        "indices_expr": _tensor_expr(gather_indices_name),\n                        "indices_values": _constant_int_list(model_ir.tensors.get(gather_indices_name, None)),\n                        "indices_shape": (\n                            model_ir.tensors[gather_indices_name].shape\n                            if gather_indices_name in model_ir.tensors\n                            else None\n                        ),\n                        "indices_shape_signature": (\n                            model_ir.tensors[gather_indices_name].shape_signature\n                            if gather_indices_name in model_ir.tensors\n                            else None\n                        ),\n                        "axis": gather_axis,\n                        "batch_dims": gather_batch_dims,\n                        "input_rank": gather_input_rank,\n                        "indices_name": gather_indices_name,\n                    }\n            if not reshape_is_lowered_onnx_flatten and not reshape_plain_singleton_axis_drop:\n                reshape_special_plan = _reshape_special_layout_plan(\n                    input_shape=reshape_input_shape,\n                    output_shape=reshape_output_semantic_shape,\n                    input_layout=reshape_input_layout,\n                    output_layout=reshape_output_layout,\n                )\n                reshape_pre_perm = _reshape_preserves_channel_last_sequence(\n                    reshape_input_shape,\n                    reshape_output_semantic_shape,\n                    reshape_input_layout,\n                )\n                if reshape_special_plan is not None and reshape_special_plan.get("pre_perm", None) is not None:\n                    reshape_pre_perm = list(reshape_special_plan["pre_perm"])\n                reshape_feature_last_target = _reshape_prefers_feature_last_for_adjx_batch_matmul(\n                    str(op.inputs[0]),\n                    str(outputs[0]),\n                )\n                if reshape_feature_last_target is not None:\n                    reshape_pre_perm = list(reshape_feature_last_target[0])\n            if (\n                reshape_pre_perm is None\n                and reshape_special_plan is None\n                and reshape_feature_last_target is None\n                and reshape_input_shape is not None\n                and reshape_output_shape is not None\n                and len(reshape_input_shape) == 2\n                and len(reshape_output_shape) == 4\n                and reshape_output_tensor is not None\n                and is_channel_last_logical_layout(\n                    normalize_logical_layout(reshape_output_tensor.logical_layout)\n                )\n                and all(int(dim) == 1 for dim in list(reshape_output_shape[1:-1]))\n                and _shape_lists_equal_relaxed(\n                    reshape_input_shape,\n                    [int(reshape_output_shape[0]), int(reshape_output_shape[-1])],\n                )\n            ):\n                reshape_channel_first_alias_shape = [\n                    int(reshape_output_shape[0]),\n                    int(reshape_output_shape[-1]),\n                    *[int(v) for v in list(reshape_output_shape[1:-1])],\n                ]\n            collapsed_reshape_shape_values: Optional[List[int]] = None\n            collapsed_reshape_input_expr: Optional[str] = None\n            upstream_reshape_index = producer_index.get(str(op.inputs[0]), None)\n            if (\n                upstream_reshape_index is not None\n                and len(consumer_index.get(str(op.inputs[0]), [])) == 1\n                and reshape_pre_perm is None\n                and reshape_special_plan is None\n                and reshape_feature_last_target is None\n                and reshape_channel_first_alias_shape is None\n                and gather_flatten_direct_spec is None\n            ):\n                upstream_reshape_op = model_ir.operators[int(upstream_reshape_index)]\n                candidate_shape_values = (\n                    _preferred_reshape_target_values(reshape_output_tensor)\n                    if reshape_output_tensor is not None\n                    else None\n                )\n                if candidate_shape_values is None and reshape_output_shape is not None:\n                    candidate_shape_values = [int(v) for v in list(reshape_output_shape)]\n                if (\n                    candidate_shape_values is not None\n                    and len(candidate_shape_values) > 0\n                    and all(int(v) > 0 for v in list(candidate_shape_values))\n                    and _can_emit_direct_torch_reshape_shape(candidate_shape_values, allow_zero=reshape_allow_zero)\n                    and _reshape_codegen_is_plain_data_only(upstream_reshape_op)\n                ):\n                    collapsed_reshape_shape_values = [int(v) for v in list(candidate_shape_values)]\n                    collapsed_reshape_input_expr = _tensor_expr(str(upstream_reshape_op.inputs[0]))\n            if collapsed_reshape_shape_values is not None and collapsed_reshape_input_expr is not None:\n                channel_first_tensor_expr_aliases.pop(str(outputs[0]), None)\n                forward_lines.append(\n                    f"{output_vars[0]} = torch.reshape({collapsed_reshape_input_expr}, {repr(collapsed_reshape_shape_values)})"\n                )\n                continue\n            if reshape_pre_perm is not None:\n                reshape_input_expr = f"{reshape_input_expr}.permute({\', \'.join(str(int(v)) for v in reshape_pre_perm)}).contiguous()"\n            shape_is_tensor_expr = False\n            if reshape_feature_last_target is not None:\n                feature_last_shape_values = [int(v) for v in list(reshape_feature_last_target[1])]\n                shape_expr = f"[int(v) for v in {repr(feature_last_shape_values)}]"\n            elif reshape_special_plan is not None and reshape_special_plan.get("reshape_shape", None) is not None:\n                shape_expr = repr([int(v) for v in list(reshape_special_plan["reshape_shape"])])\n            elif len(op.inputs) >= 2:\n                const_shape_values = _constant_int_list(model_ir.tensors.get(str(op.inputs[1]), None))\n                if const_shape_values is not None:\n                    direct_shape_values = [int(v) for v in list(const_shape_values)]\n                    preserve_feature_last_rank3_shape_expr = bool(\n                        reshape_output_tensor is not None\n                        and reshape_input_shape is not None\n                        and len(reshape_input_shape) == 2\n                        and len(list(reshape_output_tensor.shape)) == 3\n                        and is_channel_last_logical_layout(\n                            normalize_logical_layout(reshape_output_tensor.logical_layout)\n                        )\n                    )\n                    if _can_emit_direct_torch_reshape_shape(\n                        direct_shape_values,\n                        allow_zero=reshape_allow_zero,\n                    ) and all(int(v) > 0 for v in direct_shape_values):\n                        if preserve_feature_last_rank3_shape_expr:\n                            shape_expr = f"[int(v) for v in {repr(direct_shape_values)}]"\n                        else:\n                            shape_expr = repr(direct_shape_values)\n                    else:\n                        runtime_imports.add("_resolve_reshape_shape")\n                        shape_expr = (\n                            f"_resolve_reshape_shape({repr(direct_shape_values)}, "\n                            f"{reshape_input_expr}, allow_zero={reshape_allow_zero})"\n                        )\n                else:\n                    reconstructed_shape_expr = _reconstruct_shape_list_expr(str(op.inputs[1]))\n                    if reconstructed_shape_expr is not None:\n                        shape_expr = reconstructed_shape_expr\n                    else:\n                        shape_is_tensor_expr = True\n                        if (\n                            reshape_allow_zero\n                            or _reshape_shape_tensor_uses_runtime_dims(str(op.inputs[1]))\n                        ):\n                            runtime_imports.add("_shape_list")\n                            shape_expr = f"{_tensor_expr(str(op.inputs[1]))}.to(dtype=torch.int64).reshape(-1)"\n                        else:\n                            runtime_imports.add("_resolve_reshape_shape_tensor")\n                            shape_expr = (\n                                f"_resolve_reshape_shape_tensor({_tensor_expr(str(op.inputs[1]))}, "\n                                f"{reshape_input_expr}, allow_zero=False)"\n                            )\n            elif reshape_output_preferred_shape is not None:\n                preferred_shape_values = [int(v) for v in list(reshape_output_preferred_shape)]\n                if _can_emit_direct_torch_reshape_shape(\n                    preferred_shape_values,\n                    allow_zero=reshape_allow_zero,\n                ) and all(int(v) > 0 for v in preferred_shape_values):\n                    shape_expr = repr(preferred_shape_values)\n                else:\n                    runtime_imports.add("_resolve_reshape_shape")\n                    shape_expr = (\n                        f"_resolve_reshape_shape({repr(preferred_shape_values)}, "\n                        f"{reshape_input_expr}, allow_zero={reshape_allow_zero})"\n                    )\n            else:\n                runtime_imports.add("_resolve_reshape_shape")\n                raw_new_shape = op.options.get("onnxRawNewShape", op.options.get("newShape", []))\n                shape_expr = (\n                    f"_resolve_reshape_shape({repr([int(v) for v in list(raw_new_shape)])}, "\n                    f"{reshape_input_expr}, allow_zero={reshape_allow_zero})"\n                )\n            raw_new_shape = op.options.get("onnxRawNewShape", op.options.get("newShape", None))\n            reshape_is_gather_elements_axis_coord = bool(\n                str(op.inputs[0]).endswith("_gather_elements_indices_i32")\n                and str(outputs[0]).endswith("_gather_elements_axis_coord")\n                and reshape_input_shape is not None\n                and isinstance(raw_new_shape, list)\n                and [int(v) for v in list(raw_new_shape)] == [*reshape_input_shape, 1]\n            )\n            if reshape_is_gather_elements_axis_coord:\n                forward_lines.append(f"{output_vars[0]} = torch.unsqueeze({reshape_input_expr}, dim=-1)")\n            else:\n                reshape_shape_arg = f"_shape_list({shape_expr})" if shape_is_tensor_expr else shape_expr\n                direct_gather_reshape_expr = None\n                if (\n                    gather_flatten_direct_spec is not None\n                    and reshape_pre_perm is None\n                    and reshape_special_plan is None\n                    and reshape_feature_last_target is None\n                    and reshape_channel_first_alias_shape is None\n                ):\n                    direct_gather_reshape_expr = _direct_gather_reshape_expr(\n                        params_expr=str(gather_flatten_direct_spec["params_expr"]),\n                        indices_expr=str(gather_flatten_direct_spec["indices_expr"]),\n                        indices_values=cast(Optional[Sequence[int]], gather_flatten_direct_spec["indices_values"]),\n                        indices_shape=cast(Optional[Sequence[int]], gather_flatten_direct_spec["indices_shape"]),\n                        indices_shape_signature=cast(Optional[Sequence[int]], gather_flatten_direct_spec["indices_shape_signature"]),\n                        axis=int(gather_flatten_direct_spec["axis"]),\n                        batch_dims=int(gather_flatten_direct_spec["batch_dims"]),\n                        input_rank=int(gather_flatten_direct_spec["input_rank"]),\n                        indices_name=str(gather_flatten_direct_spec["indices_name"]),\n                        final_shape_values=(\n                            reshape_output_preferred_shape\n                            if reshape_output_preferred_shape is not None\n                            else reshape_output_shape\n                        ),\n                    )\n                if reshape_channel_first_alias_shape is not None:\n                    output_name = str(outputs[0])\n                    raw_output_var = _derived_local_var_name(f"{output_vars[0]}_cf")\n                    channel_first_tensor_expr_aliases[output_name] = raw_output_var\n                    forward_lines.append(\n                        f"{raw_output_var} = torch.reshape({reshape_input_expr}, {repr(reshape_channel_first_alias_shape)})"\n                    )\n                    if _can_omit_materialized_channel_last_alias(output_name):\n                        continue\n                    runtime_imports.add("_align_tensor_to_target_shape")\n                    forward_lines.append(\n                        f"{output_vars[0]} = _align_tensor_to_target_shape("\n                        f"{raw_output_var}.permute(0, 2, 3, 1).contiguous(), "\n                        f"{_target_shape_literal(output_name)})"\n                    )\n                else:\n                    channel_first_tensor_expr_aliases.pop(str(outputs[0]), None)\n                    forward_lines.append(\n                        f"{output_vars[0]} = {direct_gather_reshape_expr}"\n                        if direct_gather_reshape_expr is not None\n                        else f"{output_vars[0]} = torch.reshape({reshape_input_expr}, {reshape_shape_arg})"\n                    )\n            if reshape_special_plan is not None and reshape_special_plan.get("post_perm", None) is not None:\n                post_perm = [int(v) for v in list(reshape_special_plan["post_perm"])]\n                forward_lines.append(\n                    f"{output_vars[0]} = {output_vars[0]}.permute({\', \'.join(str(v) for v in post_perm)}).contiguous()"\n                )\n            continue\n        if _emit_native_transpose_op_for_codegen(\n            model_ir=model_ir,\n            op=op,\n            outputs=outputs,\n            output_vars=output_vars,\n            preserve_channel_last_tensor_names=preserve_channel_last_tensor_names,\n            consumer_index=consumer_index,\n            producer_index=producer_index,\n            channel_first_tensor_expr_aliases=channel_first_tensor_expr_aliases,\n            runtime_imports=runtime_imports,\n            forward_lines=forward_lines,\n            tensor_expr_fn=_tensor_expr,\n            tensor_expr_for_channel_first_bridge_fn=_tensor_expr_for_channel_first_bridge,\n            can_fold_channel_last_alias_slice_consumer_fn=_can_fold_channel_last_alias_slice_consumer,\n            all_consumers_are_channel_first_binary_ops_fn=_all_consumers_are_channel_first_binary_ops,\n            can_omit_materialized_channel_last_alias_fn=_can_omit_materialized_channel_last_alias,\n            has_channel_last_consumer_hint_for_same_shape_transpose_fn=_has_channel_last_consumer_hint_for_same_shape_transpose,\n            is_batchless_rank3_public_output_transpose_fn=_is_batchless_rank3_public_output_transpose,\n            target_shape_literal_fn=_target_shape_literal,\n        ):\n            continue\n        if _emit_native_shape_transform_misc_op_for_codegen(\n            model_ir=model_ir,\n            op=op,\n            op_type=op_type,\n            outputs=outputs,\n            output_vars=output_vars,\n            runtime_imports=runtime_imports,\n            forward_lines=forward_lines,\n            tensor_expr_fn=_tensor_expr,\n            axis_expr_from_input_fn=_axis_expr_from_input,\n        ):\n            continue\n        if _emit_native_concat_op_for_codegen(\n            model_ir=model_ir,\n            op=op,\n            op_index=op_index,\n            outputs=outputs,\n            output_vars=output_vars,\n            output_target_shape=output_target_shape,\n            channel_first_tensor_expr_aliases=channel_first_tensor_expr_aliases,\n            runtime_imports=runtime_imports,\n            forward_lines=forward_lines,\n            tensor_expr_fn=_tensor_expr,\n            derived_local_var_name_fn=_derived_local_var_name,\n            activation_lines_fn=_activation_lines,\n            resolve_concat_axis_for_channel_first_fn=_resolve_concat_axis_for_channel_first,\n            channel_first_concat_input_expr_fn=_channel_first_concat_input_expr,\n            tensor_shape_list_fn=_tensor_shape_list,\n            can_omit_materialized_channel_last_alias_fn=_can_omit_materialized_channel_last_alias,\n            target_shape_literal_fn=_target_shape_literal,\n            tensor_exact_static_shape_list_fn=_tensor_exact_static_shape_list,\n            target_shape_values_fn=_target_shape_values,\n        ):\n            continue\n        if op_type == "SLICE":\n            slice_input_name = str(op.inputs[0])\n            slice_input_tensor = model_ir.tensors.get(slice_input_name, None)\n            slice_input_expr = _tensor_expr(slice_input_name)\n            slice_begin_values = _constant_int_list(model_ir.tensors.get(str(op.inputs[1]), None)) or []\n            slice_size_values = _constant_int_list(model_ir.tensors.get(str(op.inputs[2]), None)) or []\n            slice_input_rank = len(model_ir.tensors[slice_input_name].shape)\n            slice_input_shape = model_ir.tensors[slice_input_name].shape\n            slice_output_tensor = model_ir.tensors.get(outputs[0], None)\n            slice_folded_channel_first = False\n            if (\n                slice_input_tensor is not None\n                and slice_input_name in channel_first_tensor_expr_aliases\n                and is_channel_last_logical_layout(\n                    normalize_logical_layout(slice_input_tensor.logical_layout)\n                )\n                and len(slice_begin_values) == slice_input_rank\n                and len(slice_size_values) == slice_input_rank\n            ):\n                perm_from_cf = _perm_cf_to_cl(slice_input_rank)\n                if perm_from_cf is not None:\n                    slice_input_expr = str(channel_first_tensor_expr_aliases[slice_input_name])\n                    slice_begin_values = _remap_axis_values_through_permutation(\n                        slice_begin_values,\n                        perm_from_cf,\n                    )\n                    slice_size_values = _remap_axis_values_through_permutation(\n                        slice_size_values,\n                        perm_from_cf,\n                    )\n                    perm_to_cf = _perm_cl_to_cf(slice_input_rank)\n                    if perm_to_cf is not None:\n                        slice_input_shape = _permute_shape(slice_input_shape, perm_to_cf) or slice_input_shape\n                    slice_folded_channel_first = True\n            direct_slice_expr = _direct_slice_expr(\n                x_expr=slice_input_expr,\n                begin_values=slice_begin_values,\n                size_values=slice_size_values,\n                input_rank=slice_input_rank,\n                input_shape=slice_input_shape,\n            )\n            if direct_slice_expr is not None:\n                if slice_folded_channel_first and slice_output_tensor is not None:\n                    output_name = str(outputs[0])\n                    output_rank = len(list(slice_output_tensor.shape))\n                    output_layout = normalize_logical_layout(slice_output_tensor.logical_layout)\n                    raw_output_var = (\n                        output_vars[0]\n                        if is_channel_first_logical_layout(output_layout)\n                        else _derived_local_var_name(f"{output_vars[0]}_cf")\n                    )\n                    channel_first_tensor_expr_aliases[output_name] = raw_output_var\n                    forward_lines.append(f"{raw_output_var} = {direct_slice_expr}")\n                    if raw_output_var != output_vars[0]:\n                        if _can_keep_channel_first_slice_output(output_name):\n                            continue\n                        perm_to_output = logical_layout_permutation(\n                            source_layout=channel_first_logical_layout(output_rank),\n                            target_layout=output_layout,\n                        )\n                        if perm_to_output is None:\n                            raise ModelIRPyTorchExportError(\n                                "Native PyTorch-like model.py codegen could not derive a folded slice layout bridge. "\n                                f"output={output_name} output_layout={output_layout} rank={output_rank}"\n                            )\n                        runtime_imports.add("_align_tensor_to_target_shape")\n                        forward_lines.append(\n                            f"{output_vars[0]} = _align_tensor_to_target_shape("\n                            f"{raw_output_var}.permute({\', \'.join(str(int(v)) for v in perm_to_output)}).contiguous(), "\n                            f"{_target_shape_literal(output_name)})"\n                        )\n                    continue\n                forward_lines.append(f"{output_vars[0]} = {direct_slice_expr}")\n            else:\n                if_axis0_tensor_mux_match = _match_if_axis0_tensor_mux_slice(op)\n                if if_axis0_tensor_mux_match is not None:\n                    runtime_imports.add("_apply_if_axis0_tensor_mux")\n                    forward_lines.append(\n                        f"{output_vars[0]} = _apply_if_axis0_tensor_mux("\n                        f"{_tensor_expr(if_axis0_tensor_mux_match[\'cond_name\'])}, "\n                        f"{_tensor_expr(if_axis0_tensor_mux_match[\'then_name\'])}, "\n                        f"{_tensor_expr(if_axis0_tensor_mux_match[\'else_name\'])}, "\n                        f"{_tensor_expr(str(op.inputs[0]))}, "\n                        f"{_tensor_expr(str(op.inputs[1]))}, "\n                        f"{_tensor_expr(str(op.inputs[2]))}, "\n                        f"target_shape={output_target_shape}, "\n                        "use_export_mode=self._onnx2tf_torch_export_mode)"\n                    )\n                else:\n                    runtime_imports.add("_apply_slice")\n                    forward_lines.append(\n                        f"{output_vars[0]} = _apply_slice({_tensor_expr(str(op.inputs[0]))}, {_tensor_expr(str(op.inputs[1]))}, {_tensor_expr(str(op.inputs[2]))}, target_shape={output_target_shape})"\n                    )\n            continue\n        if op_type == "STRIDED_SLICE":\n            options = dict(op.options)\n            strided_slice_input_name = str(op.inputs[0])\n            strided_slice_input_tensor = model_ir.tensors.get(strided_slice_input_name, None)\n            strided_slice_input_expr = _tensor_expr(strided_slice_input_name)\n            strided_slice_begin_values = _constant_int_list(model_ir.tensors.get(str(op.inputs[1]), None)) or []\n            strided_slice_end_values = _constant_int_list(model_ir.tensors.get(str(op.inputs[2]), None)) or []\n            strided_slice_values = _constant_int_list(model_ir.tensors.get(str(op.inputs[3]), None)) or []\n            strided_slice_input_rank = len(model_ir.tensors[strided_slice_input_name].shape)\n            strided_slice_begin_mask = int(options.get("beginMask", 0))\n            strided_slice_end_mask = int(options.get("endMask", 0))\n            strided_slice_folded_channel_first = False\n            if (\n                strided_slice_input_tensor is not None\n                and strided_slice_input_name in channel_first_tensor_expr_aliases\n                and is_channel_last_logical_layout(\n                    normalize_logical_layout(strided_slice_input_tensor.logical_layout)\n                )\n                and len(strided_slice_begin_values) == strided_slice_input_rank\n                and len(strided_slice_end_values) == strided_slice_input_rank\n                and len(strided_slice_values) == strided_slice_input_rank\n                and int(options.get("ellipsisMask", 0)) == 0\n                and int(options.get("newAxisMask", 0)) == 0\n                and int(options.get("shrinkAxisMask", 0)) == 0\n            ):\n                perm_from_cf = _perm_cf_to_cl(strided_slice_input_rank)\n                if perm_from_cf is not None:\n                    strided_slice_input_expr = str(channel_first_tensor_expr_aliases[strided_slice_input_name])\n                    strided_slice_begin_values = _remap_axis_values_through_permutation(\n                        strided_slice_begin_values,\n                        perm_from_cf,\n                    )\n                    strided_slice_end_values = _remap_axis_values_through_permutation(\n                        strided_slice_end_values,\n                        perm_from_cf,\n                    )\n                    strided_slice_values = _remap_axis_values_through_permutation(\n                        strided_slice_values,\n                        perm_from_cf,\n                    )\n                    strided_slice_begin_mask = _remap_mask_bits_through_permutation(\n                        strided_slice_begin_mask,\n                        perm_from_cf,\n                    )\n                    strided_slice_end_mask = _remap_mask_bits_through_permutation(\n                        strided_slice_end_mask,\n                        perm_from_cf,\n                    )\n                    strided_slice_folded_channel_first = True\n            direct_strided_slice_expr = _direct_strided_slice_expr(\n                x_expr=strided_slice_input_expr,\n                begin_values=strided_slice_begin_values,\n                end_values=strided_slice_end_values,\n                stride_values=strided_slice_values,\n                begin_mask=strided_slice_begin_mask,\n                end_mask=strided_slice_end_mask,\n                input_rank=strided_slice_input_rank,\n            )\n            if direct_strided_slice_expr is None:\n                direct_strided_slice_expr = _direct_symbolic_strided_slice_expr(\n                    x_expr=strided_slice_input_expr,\n                    begin_values=strided_slice_begin_values,\n                    stride_values=strided_slice_values,\n                    begin_mask=strided_slice_begin_mask,\n                    end_mask=strided_slice_end_mask,\n                    input_rank=strided_slice_input_rank,\n                    end_list_expr=_reconstruct_shape_list_expr(str(op.inputs[2])),\n                    end_scalar_expr=_reconstruct_shape_scalar_expr(str(op.inputs[2])),\n                )\n            if direct_strided_slice_expr is not None:\n                strided_slice_output_tensor = model_ir.tensors.get(outputs[0], None)\n                if strided_slice_folded_channel_first and strided_slice_output_tensor is not None:\n                    output_name = str(outputs[0])\n                    output_rank = len(list(strided_slice_output_tensor.shape))\n                    output_layout = normalize_logical_layout(strided_slice_output_tensor.logical_layout)\n                    raw_output_var = (\n                        output_vars[0]\n                        if is_channel_first_logical_layout(output_layout)\n                        else _derived_local_var_name(f"{output_vars[0]}_cf")\n                    )\n                    channel_first_tensor_expr_aliases[output_name] = raw_output_var\n                    forward_lines.append(f"{raw_output_var} = {direct_strided_slice_expr}")\n                    if raw_output_var != output_vars[0]:\n                        if _can_keep_channel_first_slice_output(output_name):\n                            continue\n                        perm_to_output = logical_layout_permutation(\n                            source_layout=channel_first_logical_layout(output_rank),\n                            target_layout=output_layout,\n                        )\n                        if perm_to_output is None:\n                            raise ModelIRPyTorchExportError(\n                                "Native PyTorch-like model.py codegen could not derive a folded strided-slice layout bridge. "\n                                f"output={output_name} output_layout={output_layout} rank={output_rank}"\n                            )\n                        runtime_imports.add("_align_tensor_to_target_shape")\n                        forward_lines.append(\n                            f"{output_vars[0]} = _align_tensor_to_target_shape("\n                            f"{raw_output_var}.permute({\', \'.join(str(int(v)) for v in perm_to_output)}).contiguous(), "\n                            f"{_target_shape_literal(output_name)})"\n                        )\n                    continue\n                forward_lines.append(f"{output_vars[0]} = {direct_strided_slice_expr}")\n            else:\n                runtime_imports.add("_apply_strided_slice")\n                forward_lines.append(\n                    f"{output_vars[0]} = _apply_strided_slice({_tensor_expr(str(op.inputs[0]))}, {_tensor_expr(str(op.inputs[1]))}, {_tensor_expr(str(op.inputs[2]))}, {_tensor_expr(str(op.inputs[3]))}, begin_mask={int(options.get(\'beginMask\', 0))}, end_mask={int(options.get(\'endMask\', 0))}, target_shape={output_target_shape})"\n                )\n            continue\n        if op_type == "SHAPE":\n            out_dtype = str(op.options.get("outType", "INT32"))\n            runtime_imports.add("_shape_tensor")\n            forward_lines.append(\n                f"{output_vars[0]} = _shape_tensor({_tensor_expr(str(op.inputs[0]))}, dtype={_torch_dtype_literal(out_dtype)}, device={_tensor_expr(str(op.inputs[0]))}.device)"\n            )\n            continue\n        if op_type == "FILL":\n            fill_value_expr = _scalar_literal_expr(str(op.inputs[1])) or f"{_tensor_expr(str(op.inputs[1]))}.reshape(-1)[0].item()"\n            fill_shape_values = _constant_int_list(model_ir.tensors.get(str(op.inputs[0]), None))\n            fill_shape_input_name = str(op.inputs[0])\n            fill_shape_input_producer = producer_by_output_name.get(fill_shape_input_name, None)\n            fill_value_tensor = model_ir.tensors.get(str(op.inputs[1]), None)\n            fill_dtype_expr = (\n                _torch_dtype_literal(str(fill_value_tensor.dtype))\n                if fill_value_tensor is not None\n                else f"{_tensor_expr(str(op.inputs[1]))}.dtype"\n            )\n            if (\n                fill_shape_input_producer is not None\n                and str(fill_shape_input_producer.op_type) == "SHAPE"\n                and len(fill_shape_input_producer.inputs) >= 1\n            ):\n                fill_device_expr = f"{_tensor_expr(str(fill_shape_input_producer.inputs[0]))}.device"\n            else:\n                fill_device_expr = f"{_tensor_expr(str(op.inputs[0]))}.device"\n            if fill_shape_values is not None:\n                fill_shape_expr = repr([int(v) for v in list(fill_shape_values)])\n                fill_shape_arg = f"[int(v) for v in {fill_shape_expr}]"\n            elif (\n                fill_shape_input_producer is not None\n                and str(fill_shape_input_producer.op_type) == "SHAPE"\n                and len(fill_shape_input_producer.inputs) >= 1\n            ):\n                runtime_imports.add("_tensor_shape_list")\n                fill_shape_arg = f"_tensor_shape_list({_tensor_expr(str(fill_shape_input_producer.inputs[0]))})"\n            else:\n                reconstructed_fill_shape_expr = _reconstruct_shape_list_expr(fill_shape_input_name)\n                if reconstructed_fill_shape_expr is not None:\n                    fill_shape_arg = reconstructed_fill_shape_expr\n                else:\n                    runtime_imports.add("_shape_list")\n                    fill_shape_arg = f"_shape_list({_tensor_expr(fill_shape_input_name)})"\n            forward_lines.append(\n                f"{output_vars[0]} = torch.full({fill_shape_arg}, {fill_value_expr}, dtype={fill_dtype_expr}, device={fill_device_expr})"\n            )\n            continue\n        if op_type == "RANDOM_STANDARD_NORMAL":\n            runtime_imports.add("_apply_random_standard_normal")\n            random_output_tensor = model_ir.tensors.get(str(outputs[0]), None)\n            random_output_shape_values = (\n                [int(v) for v in list(random_output_tensor.shape)]\n                if random_output_tensor is not None\n                and all(int(v) > 0 for v in list(random_output_tensor.shape))\n                else None\n            )\n            random_shape_values = (\n                random_output_shape_values\n                if random_output_shape_values is not None\n                else _constant_int_list(model_ir.tensors.get(str(op.inputs[0]), None))\n            )\n            if random_shape_values is not None:\n                runtime_imports.add("_module_device")\n                random_shape_expr = repr([int(v) for v in list(random_shape_values)])\n                random_device_expr = "_module_device(self)"\n            else:\n                runtime_imports.add("_shape_list")\n                random_shape_expr = _tensor_expr(str(op.inputs[0]))\n                random_device_expr = f"{_tensor_expr(str(op.inputs[0]))}.device"\n            random_dtype_expr = (\n                _torch_dtype_literal(str(random_output_tensor.dtype))\n                if random_output_tensor is not None\n                else "torch.float32"\n            )\n            random_seed_value = op.options.get("seed", None)\n            random_seed_expr = (\n                repr(int(random_seed_value))\n                if random_seed_value is not None\n                else "None"\n            )\n            forward_lines.append(\n                f"{output_vars[0]} = _apply_random_standard_normal({random_shape_expr}, dtype={random_dtype_expr}, device={random_device_expr}, seed={random_seed_expr})"\n            )\n            continue\n        if op_type == "SCATTER_ND":\n            runtime_imports.add("_apply_scatter_nd")\n            scatter_shape_values = _constant_int_list(model_ir.tensors.get(str(op.inputs[2]), None))\n            scatter_shape_expr = (\n                repr([int(v) for v in list(scatter_shape_values)])\n                if scatter_shape_values is not None\n                else _tensor_expr(str(op.inputs[2]))\n            )\n            forward_lines.append(\n                f"{output_vars[0]} = _apply_scatter_nd({_tensor_expr(str(op.inputs[0]))}, {_tensor_expr(str(op.inputs[1]))}, {scatter_shape_expr}, {output_target_shape})"\n            )\n            continue\n        if op_type == "RANGE":\n            if _range_only_feeds_identity_nms_postprocess_gathers(outputs[0]):\n                continue\n            start_tensor_expr = _tensor_expr(str(op.inputs[0]))\n            limit_tensor_expr = _tensor_expr(str(op.inputs[1]))\n            delta_tensor_expr = _tensor_expr(str(op.inputs[2]))\n            start_expr = _scalar_literal_expr(str(op.inputs[0])) or start_tensor_expr\n            limit_expr = _scalar_literal_expr(str(op.inputs[1])) or limit_tensor_expr\n            delta_expr = _scalar_literal_expr(str(op.inputs[2])) or delta_tensor_expr\n            start_value_expr = start_expr if _scalar_literal_expr(str(op.inputs[0])) is not None else f"{start_expr}.reshape(-1)[0].item()"\n            limit_value_expr = limit_expr if _scalar_literal_expr(str(op.inputs[1])) is not None else f"{limit_expr}.reshape(-1)[0].item()"\n            delta_value_expr = delta_expr if _scalar_literal_expr(str(op.inputs[2])) is not None else f"{delta_expr}.reshape(-1)[0].item()"\n            range_device_expr = (\n                f"{limit_tensor_expr}.device"\n                if _scalar_literal_expr(str(op.inputs[1])) is None else\n                f"{start_tensor_expr}.device"\n                if _scalar_literal_expr(str(op.inputs[0])) is None else\n                f"{delta_tensor_expr}.device"\n                if _scalar_literal_expr(str(op.inputs[2])) is None else\n                "_module_device(self)"\n            )\n            if (\n                _scalar_literal_expr(str(op.inputs[0])) is not None\n                and _scalar_literal_expr(str(op.inputs[1])) is not None\n                and _scalar_literal_expr(str(op.inputs[2])) is not None\n            ):\n                runtime_imports.add("_module_device")\n            range_dtype_expr = (\n                f"{limit_tensor_expr}.dtype"\n                if _scalar_literal_expr(str(op.inputs[1])) is None else\n                f"{start_tensor_expr}.dtype"\n                if _scalar_literal_expr(str(op.inputs[0])) is None else\n                f"{delta_tensor_expr}.dtype"\n                if _scalar_literal_expr(str(op.inputs[2])) is None else\n                _torch_dtype_literal(str(model_ir.tensors[str(outputs[0])].dtype))\n            )\n            forward_lines.append(\n                f"{output_vars[0]} = torch.arange(start={start_value_expr}, end={limit_value_expr}, step={delta_value_expr}, device={range_device_expr}, dtype={range_dtype_expr})"\n            )\n            continue\n        if op_type == "CUMSUM":\n            runtime_imports.add("_apply_cumsum")\n            input_name = str(op.inputs[0])\n            input_expr = _tensor_expr(input_name)\n            axis_expr = (\n                _axis_expr_from_input(str(op.inputs[1]), device_expr=input_expr)\n                if len(op.inputs) >= 2\n                else repr(int(op.options.get("axis", 0)))\n            )\n            expr = (\n                f"_apply_cumsum({input_expr}, axis={axis_expr}, "\n                f"exclusive={bool(op.options.get(\'exclusive\', False))}, "\n                f"reverse={bool(op.options.get(\'reverse\', False))})"\n            )\n            inferred_shape = _tensor_shape_list(input_name)\n            if _should_skip_align_for_shape_preserving_unary(input_name, outputs[0]):\n                forward_lines.append(f"{output_vars[0]} = {expr}")\n            else:\n                forward_lines.append(\n                    f"{output_vars[0]} = {_emit_maybe_aligned_expr(output_name=outputs[0], expr=expr, inferred_shape=inferred_shape)}"\n                )\n            continue\n        if op_type == "DEPTH_TO_SPACE":\n            block_size = int(op.options.get("blockSize", 1))\n            input_tensor = model_ir.tensors.get(str(op.inputs[0]), None)\n            output_tensor = model_ir.tensors.get(str(op.outputs[0]), None)\n            input_shape = _tensor_shape_list(str(op.inputs[0]))\n            output_shape = _tensor_shape_list(outputs[0])\n            input_layout = normalize_logical_layout(input_tensor.logical_layout) if input_tensor is not None else LOGICAL_LAYOUT_UNKNOWN\n            output_layout = normalize_logical_layout(output_tensor.logical_layout) if output_tensor is not None else LOGICAL_LAYOUT_UNKNOWN\n            use_channel_last = bool(\n                is_channel_last_logical_layout(input_layout)\n                or is_channel_last_logical_layout(output_layout)\n            )\n            inferred_channel_last = None\n            if input_shape is not None and output_shape is not None:\n                inferred_channel_last = _should_emit_channel_last_depth_to_space(\n                    input_shape=input_shape,\n                    output_shape=output_shape,\n                    block_size=block_size,\n                )\n            if inferred_channel_last is True:\n                use_channel_last = True\n            elif inferred_channel_last is False:\n                use_channel_last = False\n            input_expr = _tensor_expr(str(op.inputs[0]))\n            if use_channel_last:\n                forward_lines.append(f"_depth_to_space_x_{op_index} = {input_expr}")\n                forward_lines.append(\n                    f"_depth_to_space_n_{op_index}, _depth_to_space_h_{op_index}, _depth_to_space_w_{op_index}, _depth_to_space_c_{op_index} = _depth_to_space_x_{op_index}.shape"\n                )\n                forward_lines.append(\n                    f"{output_vars[0]} = _depth_to_space_x_{op_index}.reshape(_depth_to_space_n_{op_index}, _depth_to_space_h_{op_index}, _depth_to_space_w_{op_index}, {block_size}, {block_size}, _depth_to_space_c_{op_index} // {block_size * block_size}).permute(0, 1, 3, 2, 4, 5).reshape(_depth_to_space_n_{op_index}, _depth_to_space_h_{op_index} * {block_size}, _depth_to_space_w_{op_index} * {block_size}, _depth_to_space_c_{op_index} // {block_size * block_size})"\n                )\n            else:\n                forward_lines.append(\n                    f"{output_vars[0]} = {_emit_maybe_aligned_expr(output_name=outputs[0], expr=f\'F.pixel_shuffle({input_expr}, {block_size})\', inferred_shape=_tensor_shape_list(outputs[0]))}"\n                )\n            continue\n        if op_type == "SPACE_TO_DEPTH":\n            block_size = int(op.options.get("blockSize", 1))\n            input_expr = _tensor_expr(str(op.inputs[0]))\n            input_tensor = model_ir.tensors.get(str(op.inputs[0]), None)\n            output_tensor = model_ir.tensors.get(str(op.outputs[0]), None)\n            input_shape = _tensor_shape_list(str(op.inputs[0]))\n            output_shape = _tensor_shape_list(outputs[0])\n            input_layout = normalize_logical_layout(input_tensor.logical_layout) if input_tensor is not None else LOGICAL_LAYOUT_UNKNOWN\n            output_layout = normalize_logical_layout(output_tensor.logical_layout) if output_tensor is not None else LOGICAL_LAYOUT_UNKNOWN\n            use_channel_last = bool(\n                is_channel_last_logical_layout(input_layout)\n                or is_channel_last_logical_layout(output_layout)\n            )\n            inferred_channel_last = None\n            if input_shape is not None and output_shape is not None:\n                inferred_channel_last = _should_emit_channel_last_space_to_depth(\n                    input_shape=input_shape,\n                    output_shape=output_shape,\n                    block_size=block_size,\n                )\n            if inferred_channel_last is True:\n                use_channel_last = True\n            elif inferred_channel_last is False:\n                use_channel_last = False\n            forward_lines.append(f"_space_to_depth_x_{op_index} = {input_expr}")\n            if use_channel_last:\n                forward_lines.append(\n                    f"_space_to_depth_n_{op_index}, _space_to_depth_h_{op_index}, _space_to_depth_w_{op_index}, _space_to_depth_c_{op_index} = _space_to_depth_x_{op_index}.shape"\n                )\n                forward_lines.append(\n                    f"{output_vars[0]} = _space_to_depth_x_{op_index}.reshape(_space_to_depth_n_{op_index}, _space_to_depth_h_{op_index} // {block_size}, {block_size}, _space_to_depth_w_{op_index} // {block_size}, {block_size}, _space_to_depth_c_{op_index}).permute(0, 1, 3, 2, 4, 5).reshape(_space_to_depth_n_{op_index}, _space_to_depth_h_{op_index} // {block_size}, _space_to_depth_w_{op_index} // {block_size}, _space_to_depth_c_{op_index} * {block_size * block_size})"\n                )\n            else:\n                forward_lines.append(\n                    f"_space_to_depth_n_{op_index}, _space_to_depth_c_{op_index}, _space_to_depth_h_{op_index}, _space_to_depth_w_{op_index} = _space_to_depth_x_{op_index}.shape"\n                )\n                forward_lines.append(\n                    f"{output_vars[0]} = _space_to_depth_x_{op_index}.reshape(_space_to_depth_n_{op_index}, _space_to_depth_c_{op_index}, _space_to_depth_h_{op_index} // {block_size}, {block_size}, _space_to_depth_w_{op_index} // {block_size}, {block_size}).permute(0, 1, 3, 5, 2, 4).reshape(_space_to_depth_n_{op_index}, _space_to_depth_c_{op_index} * {block_size * block_size}, _space_to_depth_h_{op_index} // {block_size}, _space_to_depth_w_{op_index} // {block_size})"\n                )\n            continue\n        if op_type == "SOFTMAX":\n            runtime_imports.add("_apply_softmax")\n            axis = op.options.get("axis", None)\n            if axis is None and len(op.inputs) > 0:\n                input_tensor = model_ir.tensors.get(str(op.inputs[0]), None)\n                if input_tensor is not None:\n                    input_layout = normalize_logical_layout(input_tensor.logical_layout)\n                    if is_channel_first_logical_layout(input_layout) and len(list(input_tensor.shape)) >= 2:\n                        axis = 1\n            axis_expr = repr(int(axis)) if axis is not None else "None"\n            beta = float(op.options.get("beta", 1.0))\n            forward_lines.append(\n                f"{output_vars[0]} = _apply_softmax({_tensor_expr(str(op.inputs[0]))}, axis={axis_expr}, beta={beta}, target_shape={output_target_shape})"\n            )\n            continue\n        if op_type in {"ARG_MAX", "ARG_MIN"}:\n            runtime_imports.add("_normalize_dim")\n            input_expr = _tensor_expr(str(op.inputs[0]))\n            if len(op.inputs) >= 2:\n                axis_expr = _axis_expr_from_input(str(op.inputs[1]), device_expr=input_expr)\n            else:\n                axis_expr = repr(int(op.options.get("axis", 0)))\n            input_tensor = model_ir.tensors.get(str(op.inputs[0]), None)\n            output_tensor = model_ir.tensors.get(outputs[0], None)\n            keep_dims = bool(op.options.get("keepDims", False))\n            if input_tensor is not None and output_tensor is not None:\n                keep_dims = len(list(output_tensor.shape)) == len(list(input_tensor.shape))\n            reduce_fn = "torch.argmax" if op_type == "ARG_MAX" else "torch.argmin"\n            output_dtype = "INT64" if output_tensor is None else str(output_tensor.dtype)\n            forward_lines.append(\n                f"{output_vars[0]} = {reduce_fn}({input_expr}, dim=_normalize_dim({axis_expr}, {input_expr}.ndim), keepdim={keep_dims}).to(dtype={_torch_dtype_literal(output_dtype)})"\n            )\n            continue\n        if op_type == "TOPK_V2":\n            runtime_imports.add("_normalize_dim")\n            input_expr = _tensor_expr(str(op.inputs[0]))\n            topk_pre_perm, topk_index_post_perm = _topk_codegen_layout_bridge(\n                input_name=str(op.inputs[0]),\n                value_output_name=str(outputs[0]),\n                index_output_name=str(outputs[1]) if len(outputs) > 1 else None,\n            )\n            if topk_pre_perm is not None:\n                input_expr = f"{input_expr}.permute({\', \'.join(str(int(v)) for v in topk_pre_perm)}).contiguous()"\n            input_tensor = model_ir.tensors.get(str(op.inputs[0]), None)\n            output_value_tensor = model_ir.tensors.get(outputs[0], None)\n            output_index_tensor = model_ir.tensors.get(outputs[1], None) if len(outputs) > 1 else None\n            k_literal = _int_scalar_literal_expr(str(op.inputs[1]))\n            k_reconstructed_scalar_expr = None if k_literal is not None else _reconstruct_shape_scalar_expr(str(op.inputs[1]))\n            k_expr = (\n                repr(int(k_literal))\n                if k_literal is not None else\n                (\n                    f"{k_reconstructed_scalar_expr}"\n                    if k_reconstructed_scalar_expr is not None else\n                f"int({_tensor_expr(str(op.inputs[1]))}.reshape(-1)[0].to(dtype=torch.int64).item())"\n                )\n            )\n            axis_expr = repr(int(op.options.get("axis", -1)))\n            largest = bool(op.options.get("largest", True))\n            sorted_output = bool(op.options.get("sorted", True))\n            output_value_shape_signature = (\n                [int(v) for v in list(output_value_tensor.shape_signature)]\n                if output_value_tensor is not None and output_value_tensor.shape_signature is not None\n                else None\n            )\n            input_shape_signature = (\n                [int(v) for v in list(input_tensor.shape_signature)]\n                if input_tensor is not None and input_tensor.shape_signature is not None\n                else None\n            )\n            can_emit_full_sort = (\n                k_literal is None\n                and\n                output_value_shape_signature is not None\n                and input_shape_signature is not None\n                and output_value_shape_signature == input_shape_signature\n            )\n            if can_emit_full_sort:\n                forward_lines.append(\n                    f"{output_vars[0]}, {output_vars[1]} = torch.sort({input_expr}, dim=_normalize_dim({axis_expr}, {input_expr}.ndim), descending={largest})"\n                )\n            else:\n                forward_lines.append(\n                    f"{output_vars[0]}, {output_vars[1]} = torch.topk({input_expr}, k={k_expr}, dim=_normalize_dim({axis_expr}, {input_expr}.ndim), largest={largest}, sorted={sorted_output})"\n                )\n            runtime_shape_uncertain_tensors.update(outputs)\n            if topk_index_post_perm is not None:\n                forward_lines.append(\n                    f"{output_vars[1]} = {output_vars[1]}.permute({\', \'.join(str(int(v)) for v in topk_index_post_perm)}).contiguous()"\n                )\n            if output_value_tensor is not None:\n                forward_lines.append(\n                    f"{output_vars[0]} = {output_vars[0]}.to(dtype={_torch_dtype_literal(str(output_value_tensor.dtype))})"\n                )\n            if output_index_tensor is not None:\n                forward_lines.append(\n                    f"{output_vars[1]} = {output_vars[1]}.to(dtype={_torch_dtype_literal(str(output_index_tensor.dtype))})"\n                )\n            continue\n        if op_type == "AVERAGE_POOL_2D":\n            options = dict(op.options)\n            pool_input_tensor = model_ir.tensors.get(str(op.inputs[0]), None)\n            pool_output_tensor = model_ir.tensors.get(str(op.outputs[0]), None) if len(op.outputs) == 1 else None\n            effective_layout = _infer_effective_rank4_runtime_layout(str(op.inputs[0]))\n            pool_shape_implies_channel_last: Optional[bool] = None\n            if pool_input_tensor is not None and pool_output_tensor is not None:\n                input_shape = [int(v) for v in list(pool_input_tensor.shape)]\n                output_shape = [int(v) for v in list(pool_output_tensor.shape)]\n                if len(input_shape) == 4 and len(output_shape) == 4:\n                    if input_shape[3] == output_shape[1] and input_shape[1] != output_shape[1]:\n                        pool_shape_implies_channel_last = True\n                    elif input_shape[1] == output_shape[1] and input_shape[3] != output_shape[1]:\n                        pool_shape_implies_channel_last = False\n            pool_alias_expr = channel_first_tensor_expr_aliases.get(str(op.inputs[0]), None)\n            pool_distinct_channel_first_alias = (\n                pool_alias_expr is not None\n                and pool_alias_expr != _tensor_expr(str(op.inputs[0]))\n            )\n            pool_uses_channel_first_alias = (\n                pool_distinct_channel_first_alias\n            )\n            pool_input_expr = (\n                str(pool_alias_expr)\n                if pool_uses_channel_first_alias else\n                _tensor_expr(str(op.inputs[0]))\n            )\n            pool_as_channel_last_expr = (\n                "False"\n                if pool_uses_channel_first_alias else\n                "True" if pool_shape_implies_channel_last is True else\n                "False" if pool_shape_implies_channel_last is False else\n                "True" if effective_layout == "NHWC" else\n                "False" if effective_layout == "NCHW" else\n                "None"\n            )\n            runtime_imports.add("_apply_pool2d")\n            forward_lines.append(\n                f"{output_vars[0]} = _apply_pool2d({pool_input_expr}, "\n                f"filter_height={int(options.get(\'filterHeight\', 1))}, "\n                f"filter_width={int(options.get(\'filterWidth\', 1))}, "\n                f"stride_h={int(options.get(\'strideH\', 1))}, "\n                f"stride_w={int(options.get(\'strideW\', 1))}, "\n                f"padding={str(options.get(\'padding\', \'VALID\')).upper()!r}, "\n                f"target_shape={output_target_shape}, "\n                f"is_max_pool=False, channel_last={pool_as_channel_last_expr})"\n            )\n            forward_lines.extend(_activation_lines(output_vars[0], str(options.get("fusedActivationFunction", "NONE"))))\n            continue\n        if op_type == "MAX_POOL_2D":\n            options = dict(op.options)\n            pool_input_tensor = model_ir.tensors.get(str(op.inputs[0]), None)\n            pool_output_tensor = model_ir.tensors.get(str(op.outputs[0]), None) if len(op.outputs) == 1 else None\n            effective_layout = _infer_effective_rank4_runtime_layout(str(op.inputs[0]))\n            pool_shape_implies_channel_last: Optional[bool] = None\n            if pool_input_tensor is not None and pool_output_tensor is not None:\n                input_shape = [int(v) for v in list(pool_input_tensor.shape)]\n                output_shape = [int(v) for v in list(pool_output_tensor.shape)]\n                if len(input_shape) == 4 and len(output_shape) == 4:\n                    if input_shape[3] == output_shape[1] and input_shape[1] != output_shape[1]:\n                        pool_shape_implies_channel_last = True\n                    elif input_shape[1] == output_shape[1] and input_shape[3] != output_shape[1]:\n                        pool_shape_implies_channel_last = False\n            pool_alias_expr = channel_first_tensor_expr_aliases.get(str(op.inputs[0]), None)\n            pool_distinct_channel_first_alias = (\n                pool_alias_expr is not None\n                and pool_alias_expr != _tensor_expr(str(op.inputs[0]))\n            )\n            pool_uses_channel_first_alias = (\n                pool_distinct_channel_first_alias\n            )\n            pool_input_expr = (\n                str(pool_alias_expr)\n                if pool_uses_channel_first_alias else\n                _tensor_expr(str(op.inputs[0]))\n            )\n            pool_as_channel_last_expr = (\n                "False"\n                if pool_uses_channel_first_alias else\n                "True" if pool_shape_implies_channel_last is True else\n                "False" if pool_shape_implies_channel_last is False else\n                "True" if effective_layout == "NHWC" else\n                "False" if effective_layout == "NCHW" else\n                "None"\n            )\n            runtime_imports.add("_apply_pool2d")\n            forward_lines.append(\n                f"{output_vars[0]} = _apply_pool2d({pool_input_expr}, "\n                f"filter_height={int(options.get(\'filterHeight\', 1))}, "\n                f"filter_width={int(options.get(\'filterWidth\', 1))}, "\n                f"stride_h={int(options.get(\'strideH\', 1))}, "\n                f"stride_w={int(options.get(\'strideW\', 1))}, "\n                f"padding={str(options.get(\'padding\', \'VALID\')).upper()!r}, "\n                f"target_shape={output_target_shape}, "\n                f"is_max_pool=True, channel_last={pool_as_channel_last_expr})"\n            )\n            forward_lines.extend(_activation_lines(output_vars[0], str(options.get("fusedActivationFunction", "NONE"))))\n            continue\n        if op_type == "RESIZE_NEAREST_NEIGHBOR":\n            size_tensor = model_ir.tensors.get(str(op.inputs[1]), None)\n            size_literal = _python_literal_for_constant_tensor(size_tensor) if size_tensor is not None else None\n            size_values = _constant_int_list(size_tensor) if size_tensor is not None else None\n            resize_target_shape = _resize_target_shape_literal(outputs[0], str(op.inputs[0]))\n            resize_input_tensor = model_ir.tensors.get(str(op.inputs[0]), None)\n            resize_output_tensor = model_ir.tensors.get(str(outputs[0]), None) if len(outputs) == 1 else None\n            resize_effective_layout = _infer_effective_rank4_runtime_layout(str(op.inputs[0]))\n            resize_shape_implies_channel_last: Optional[bool] = None\n            if resize_input_tensor is not None and resize_output_tensor is not None:\n                resize_input_shape = [int(v) for v in list(resize_input_tensor.shape)]\n                resize_output_shape = [int(v) for v in list(resize_output_tensor.shape)]\n                if len(resize_input_shape) == 4 and len(resize_output_shape) == 4:\n                    if resize_input_shape[3] == resize_output_shape[3] and resize_input_shape[1] != resize_output_shape[3]:\n                        resize_shape_implies_channel_last = True\n                    elif resize_input_shape[1] == resize_output_shape[1] and resize_input_shape[3] != resize_output_shape[1]:\n                        resize_shape_implies_channel_last = False\n            resize_alias_expr = channel_first_tensor_expr_aliases.get(str(op.inputs[0]), None)\n            resize_distinct_channel_first_alias = (\n                resize_alias_expr is not None\n                and resize_alias_expr != _tensor_expr(str(op.inputs[0]))\n            )\n            resize_uses_channel_first_alias = (\n                resize_distinct_channel_first_alias\n            )\n            resize_input_expr = (\n                str(resize_alias_expr)\n                if resize_uses_channel_first_alias else\n                _tensor_expr(str(op.inputs[0]))\n            )\n            output_cf_shape = _rank4_channel_first_shape_for_tensor(outputs[0])\n            can_emit_resize_cf_alias = (\n                resize_uses_channel_first_alias\n                and output_cf_shape is not None\n                and len(output_cf_shape) == 4\n                and all(int(dim) > 0 for dim in output_cf_shape)\n            )\n            resize_channel_last_expr = (\n                "False"\n                if resize_uses_channel_first_alias else\n                "True" if resize_shape_implies_channel_last is True else\n                "False" if resize_shape_implies_channel_last is False else\n                "True" if resize_effective_layout == "NHWC" else\n                "False" if resize_effective_layout == "NCHW" else\n                "None"\n            )\n            size_expr = (\n                f"{size_literal}"\n                if size_literal is not None else\n                _tensor_expr(str(op.inputs[1]))\n            )\n            emitted_direct_resize = False\n            input_cf_shape = _rank4_channel_first_shape_for_tensor(str(op.inputs[0]))\n            output_shape = _tensor_shape_list(outputs[0])\n            if resize_channel_last_expr == "False" and input_cf_shape is not None and size_values is not None and len(size_values) == 2:\n                raw_resize_shape = [int(input_cf_shape[0]), int(input_cf_shape[1]), int(size_values[0]), int(size_values[1])]\n                resize_expr = f"F.interpolate({resize_input_expr}, size={repr([int(size_values[0]), int(size_values[1])])}, mode=\'nearest\')"\n                if output_shape == raw_resize_shape:\n                    channel_first_tensor_expr_aliases.pop(str(outputs[0]), None)\n                    forward_lines.append(f"{output_vars[0]} = {resize_expr}")\n                    emitted_direct_resize = True\n                elif output_shape == [int(raw_resize_shape[0]), int(raw_resize_shape[2]), int(raw_resize_shape[3]), int(raw_resize_shape[1])]:\n                    raw_output_var = _derived_local_var_name(f"{output_vars[0]}_cf")\n                    channel_first_tensor_expr_aliases[str(outputs[0])] = raw_output_var\n                    forward_lines.append(f"{raw_output_var} = {resize_expr}")\n                    if not _can_omit_materialized_channel_last_alias(outputs[0]):\n                        runtime_imports.add("_align_tensor_to_target_shape")\n                        forward_lines.append(\n                            f"{output_vars[0]} = _align_tensor_to_target_shape({raw_output_var}.permute(0, 2, 3, 1).contiguous(), {resize_target_shape})"\n                        )\n                    emitted_direct_resize = True\n            if not emitted_direct_resize:\n                runtime_imports.add("_apply_resize")\n                if can_emit_resize_cf_alias:\n                    raw_output_var = _derived_local_var_name(f"{output_vars[0]}_cf")\n                    channel_first_tensor_expr_aliases[str(outputs[0])] = raw_output_var\n                    forward_lines.append(\n                        f"{raw_output_var} = _apply_resize({resize_input_expr}, {size_expr}, method=\'nearest\', target_shape={repr(output_cf_shape)}, channel_last=False)"\n                    )\n                    if not _can_omit_materialized_channel_last_alias(outputs[0]):\n                        runtime_imports.add("_align_tensor_to_target_shape")\n                        forward_lines.append(\n                            f"{output_vars[0]} = _align_tensor_to_target_shape({raw_output_var}.permute(0, 2, 3, 1).contiguous(), {resize_target_shape})"\n                        )\n                else:\n                    channel_first_tensor_expr_aliases.pop(str(outputs[0]), None)\n                    forward_lines.append(\n                        f"{output_vars[0]} = _apply_resize({resize_input_expr}, {size_expr}, method=\'nearest\', target_shape={resize_target_shape}, channel_last={resize_channel_last_expr})"\n                    )\n            continue\n        if op_type == "RESIZE_BILINEAR":\n            size_tensor = model_ir.tensors.get(str(op.inputs[1]), None)\n            size_literal = _python_literal_for_constant_tensor(size_tensor) if size_tensor is not None else None\n            size_values = _constant_int_list(size_tensor) if size_tensor is not None else None\n            align_corners = bool(op.options.get("alignCorners", False))\n            half_pixel_centers = bool(op.options.get("halfPixelCenters", False))\n            resize_target_shape = _resize_target_shape_literal(outputs[0], str(op.inputs[0]))\n            resize_input_tensor = model_ir.tensors.get(str(op.inputs[0]), None)\n            resize_output_tensor = model_ir.tensors.get(str(outputs[0]), None) if len(outputs) == 1 else None\n            resize_effective_layout = _infer_effective_rank4_runtime_layout(str(op.inputs[0]))\n            resize_shape_implies_channel_last: Optional[bool] = None\n            if resize_input_tensor is not None and resize_output_tensor is not None:\n                resize_input_shape = [int(v) for v in list(resize_input_tensor.shape)]\n                resize_output_shape = [int(v) for v in list(resize_output_tensor.shape)]\n                if len(resize_input_shape) == 4 and len(resize_output_shape) == 4:\n                    if resize_input_shape[3] == resize_output_shape[3] and resize_input_shape[1] != resize_output_shape[3]:\n                        resize_shape_implies_channel_last = True\n                    elif resize_input_shape[1] == resize_output_shape[1] and resize_input_shape[3] != resize_output_shape[1]:\n                        resize_shape_implies_channel_last = False\n            resize_alias_expr = channel_first_tensor_expr_aliases.get(str(op.inputs[0]), None)\n            resize_distinct_channel_first_alias = (\n                resize_alias_expr is not None\n                and resize_alias_expr != _tensor_expr(str(op.inputs[0]))\n            )\n            resize_uses_channel_first_alias = (\n                resize_distinct_channel_first_alias\n            )\n            resize_input_expr = (\n                str(resize_alias_expr)\n                if resize_uses_channel_first_alias else\n                _tensor_expr(str(op.inputs[0]))\n            )\n            output_cf_shape = _rank4_channel_first_shape_for_tensor(outputs[0])\n            can_emit_resize_cf_alias = (\n                resize_uses_channel_first_alias\n                and output_cf_shape is not None\n                and len(output_cf_shape) == 4\n                and all(int(dim) > 0 for dim in output_cf_shape)\n            )\n            resize_channel_last_expr = (\n                "False"\n                if resize_uses_channel_first_alias else\n                "True" if resize_shape_implies_channel_last is True else\n                "False" if resize_shape_implies_channel_last is False else\n                "True" if resize_effective_layout == "NHWC" else\n                "False" if resize_effective_layout == "NCHW" else\n                "None"\n            )\n            size_expr = (\n                f"{size_literal}"\n                if size_literal is not None else\n                _tensor_expr(str(op.inputs[1]))\n            )\n            emitted_direct_resize = False\n            input_cf_shape = _rank4_channel_first_shape_for_tensor(str(op.inputs[0]))\n            output_shape = _tensor_shape_list(outputs[0])\n            if (\n                resize_channel_last_expr == "False"\n                and not half_pixel_centers\n                and input_cf_shape is not None\n                and size_values is not None\n                and len(size_values) == 2\n            ):\n                raw_resize_shape = [int(input_cf_shape[0]), int(input_cf_shape[1]), int(size_values[0]), int(size_values[1])]\n                resize_expr = (\n                    f"F.interpolate({resize_input_expr}, size={repr([int(size_values[0]), int(size_values[1])])}, "\n                    f"mode=\'bilinear\', align_corners={align_corners})"\n                )\n                if output_shape == raw_resize_shape:\n                    channel_first_tensor_expr_aliases.pop(str(outputs[0]), None)\n                    forward_lines.append(f"{output_vars[0]} = {resize_expr}")\n                    emitted_direct_resize = True\n                elif output_shape == [int(raw_resize_shape[0]), int(raw_resize_shape[2]), int(raw_resize_shape[3]), int(raw_resize_shape[1])]:\n                    raw_output_var = _derived_local_var_name(f"{output_vars[0]}_cf")\n                    channel_first_tensor_expr_aliases[str(outputs[0])] = raw_output_var\n                    forward_lines.append(f"{raw_output_var} = {resize_expr}")\n                    if not _can_omit_materialized_channel_last_alias(outputs[0]):\n                        runtime_imports.add("_align_tensor_to_target_shape")\n                        forward_lines.append(\n                            f"{output_vars[0]} = _align_tensor_to_target_shape({raw_output_var}.permute(0, 2, 3, 1).contiguous(), {resize_target_shape})"\n                        )\n                    emitted_direct_resize = True\n            if not emitted_direct_resize:\n                runtime_imports.add("_apply_resize")\n                if can_emit_resize_cf_alias:\n                    raw_output_var = _derived_local_var_name(f"{output_vars[0]}_cf")\n                    channel_first_tensor_expr_aliases[str(outputs[0])] = raw_output_var\n                    forward_lines.append(\n                        f"{raw_output_var} = _apply_resize({resize_input_expr}, {size_expr}, method=\'bilinear\', target_shape={repr(output_cf_shape)}, align_corners={align_corners}, half_pixel_centers={half_pixel_centers}, channel_last=False)"\n                    )\n                    if not _can_omit_materialized_channel_last_alias(outputs[0]):\n                        runtime_imports.add("_align_tensor_to_target_shape")\n                        forward_lines.append(\n                            f"{output_vars[0]} = _align_tensor_to_target_shape({raw_output_var}.permute(0, 2, 3, 1).contiguous(), {resize_target_shape})"\n                        )\n                else:\n                    channel_first_tensor_expr_aliases.pop(str(outputs[0]), None)\n                    forward_lines.append(\n                        f"{output_vars[0]} = _apply_resize({resize_input_expr}, {size_expr}, method=\'bilinear\', target_shape={resize_target_shape}, align_corners={align_corners}, half_pixel_centers={half_pixel_centers}, channel_last={resize_channel_last_expr})"\n                    )\n            continue\n        if op_type in {"SUM", "MEAN", "REDUCE_MAX", "REDUCE_MIN", "REDUCE_PROD", "REDUCE_ANY"}:\n            runtime_imports.update({"_normalize_axes"})\n            reducer_map = {\n                "SUM": "_reduce_sum",\n                "MEAN": "_reduce_mean",\n                "REDUCE_MAX": "_reduce_max",\n                "REDUCE_MIN": "_reduce_min",\n                "REDUCE_PROD": "_reduce_prod",\n                "REDUCE_ANY": "_reduce_any",\n            }\n            runtime_imports.add(reducer_map[op_type])\n            axis_values = None\n            if len(op.inputs) >= 2:\n                axis_values = _constant_int_list(model_ir.tensors.get(str(op.inputs[1]), None))\n            axis_expr = (\n                f"_normalize_axes({repr(axis_values)}, {_tensor_expr(str(op.inputs[0]))}.ndim)"\n                if axis_values is not None\n                else f"_normalize_axes({_tensor_expr(str(op.inputs[1]))}, {_tensor_expr(str(op.inputs[0]))}.ndim)"\n                if len(op.inputs) >= 2\n                else "None"\n            )\n            keepdims = bool(op.options.get("keepDims", True))\n            reduction_output_name = str(outputs[0])\n            reduction_output_tensor = model_ir.tensors.get(reduction_output_name, None)\n            reduction_output_rank = len(list(reduction_output_tensor.shape)) if reduction_output_tensor is not None else 0\n            reduction_output_layout = normalize_logical_layout(\n                reduction_output_tensor.logical_layout if reduction_output_tensor is not None else LOGICAL_LAYOUT_UNKNOWN\n            )\n            reduction_raw_output_var = output_vars[0]\n            reduction_input_name = str(op.inputs[0])\n            reduction_input_alias = channel_first_tensor_expr_aliases.get(reduction_input_name, None)\n            reduction_input_expr = _tensor_expr(reduction_input_name)\n            reduction_input_var = tensor_var_names.get(reduction_input_name, reduction_input_expr)\n            reduction_uses_channel_first_alias = (\n                reduction_input_alias is not None\n                or reduction_input_expr != reduction_input_var\n            )\n            reduction_shape_implies_channel_last = (\n                reduction_output_tensor is not None\n                and reduction_output_rank == 4\n                and int(reduction_output_tensor.shape[1]) == 1\n                and int(reduction_output_tensor.shape[2]) == 1\n                and int(reduction_output_tensor.shape[3]) > 1\n            )\n            reduction_input_tensor = model_ir.tensors.get(reduction_input_name, None)\n            reduction_input_is_ambiguous_channel_last_gap = (\n                reduction_input_tensor is not None\n                and len(list(reduction_input_tensor.shape)) == 4\n                and is_channel_last_logical_layout(\n                    normalize_logical_layout(reduction_input_tensor.logical_layout)\n                )\n                and axis_values is not None\n                and [int(v) for v in list(axis_values)] == [1, 2]\n                and len(\n                    set(\n                        int(dim)\n                        for dim in list(reduction_input_tensor.shape)[1:]\n                        if int(dim) > 0\n                    )\n                )\n                != len(\n                    [\n                        int(dim)\n                        for dim in list(reduction_input_tensor.shape)[1:]\n                        if int(dim) > 0\n                    ]\n                )\n            )\n            if (\n                op_type == "MEAN"\n                and keepdims\n                and reduction_input_is_ambiguous_channel_last_gap\n                and reduction_input_alias is not None\n            ):\n                inferred_reduction_shape = _infer_reduction_shape(\n                    _tensor_shape_list(str(op.inputs[0])),\n                    axis_values,\n                    keepdims=keepdims,\n                )\n                reduced_expr = _emit_maybe_aligned_expr(\n                    output_name=outputs[0],\n                    expr=(\n                        f"torch.mean({reduction_input_alias}.permute(0, 2, 3, 1).contiguous(), "\n                        f"dim=[1, 2], keepdim=True)"\n                    ),\n                    inferred_shape=inferred_reduction_shape,\n                )\n                channel_first_tensor_expr_aliases.pop(reduction_output_name, None)\n                forward_lines.append(f"{output_vars[0]} = {reduced_expr}")\n                continue\n            reduction_is_channel_last_gap = (\n                op_type == "MEAN"\n                and keepdims\n                and axis_values is not None\n                and [int(v) for v in list(axis_values)] == [1, 2]\n                and reduction_shape_implies_channel_last\n                and not reduction_input_is_ambiguous_channel_last_gap\n            )\n            if reduction_is_channel_last_gap:\n                channel_first_tensor_expr_aliases[reduction_output_name] = output_vars[0]\n            if reduction_is_channel_last_gap:\n                reduction_input_expr = _tensor_expr(str(op.inputs[0]))\n                if reduction_input_expr != tensor_var_names.get(str(op.inputs[0]), reduction_input_expr):\n                    runtime_imports.add("_align_tensor_to_target_shape")\n                    forward_lines.append(\n                        f"{output_vars[0]} = _align_tensor_to_target_shape("\n                        f"torch.mean({reduction_input_expr}, dim=[2, 3], keepdim=True).permute(0, 2, 3, 1).contiguous(), "\n                        f"{_target_shape_literal(reduction_output_name)})"\n                    )\n                    channel_first_tensor_expr_aliases.pop(reduction_output_name, None)\n                    continue\n            if (\n                (reduction_uses_channel_first_alias or reduction_is_channel_last_gap)\n                and reduction_output_tensor is not None\n                and (\n                    is_channel_last_logical_layout(reduction_output_layout)\n                    or reduction_shape_implies_channel_last\n                )\n                and reduction_output_rank in {3, 4, 5}\n            ):\n                reduction_raw_output_var = _derived_local_var_name(f"{output_vars[0]}_cf")\n                channel_first_tensor_expr_aliases[reduction_output_name] = reduction_raw_output_var\n            else:\n                channel_first_tensor_expr_aliases.pop(reduction_output_name, None)\n            channel_first_reduction_plan = _channel_first_reduction_plan(op, str(op.inputs[0]))\n            if channel_first_reduction_plan is not None:\n                input_expr, channel_first_axes = channel_first_reduction_plan\n                if op_type == "MEAN":\n                    direct_expr = _direct_mean_reduction_expr(\n                        input_expr=input_expr,\n                        axes=channel_first_axes,\n                        input_rank=len(list(model_ir.tensors[str(op.inputs[0])].shape)),\n                        keepdims=keepdims,\n                    )\n                    if direct_expr is not None:\n                        reduced_expr = _emit_maybe_aligned_expr(\n                            output_name=outputs[0],\n                            expr=direct_expr,\n                            inferred_shape=_infer_reduction_shape(_tensor_shape_list(str(op.inputs[0])), axis_values, keepdims=keepdims),\n                        )\n                        forward_lines.append(f"{reduction_raw_output_var} = {reduced_expr}")\n                        if reduction_raw_output_var != output_vars[0] and not _can_omit_materialized_channel_last_alias(reduction_output_name):\n                            runtime_imports.add("_align_tensor_to_target_shape")\n                            forward_lines.append(\n                                f"{output_vars[0]} = _align_tensor_to_target_shape("\n                                f"{reduction_raw_output_var}.permute(0, 2, 3, 1).contiguous(), "\n                                f"{_target_shape_literal(reduction_output_name)})"\n                            )\n                        continue\n                axis_expr = f"_normalize_axes({repr(channel_first_axes)}, {input_expr}.ndim)"\n                reduced_expr = _emit_maybe_aligned_expr(\n                    output_name=outputs[0],\n                    expr=f"{reducer_map[op_type]}({input_expr}, {axis_expr}, {keepdims})",\n                    inferred_shape=_infer_reduction_shape(_tensor_shape_list(str(op.inputs[0])), axis_values, keepdims=keepdims),\n                )\n                forward_lines.append(f"{reduction_raw_output_var} = {reduced_expr}")\n                if reduction_raw_output_var != output_vars[0] and not _can_omit_materialized_channel_last_alias(reduction_output_name):\n                    runtime_imports.add("_align_tensor_to_target_shape")\n                    forward_lines.append(\n                        f"{output_vars[0]} = _align_tensor_to_target_shape("\n                        f"{reduction_raw_output_var}.permute(0, 2, 3, 1).contiguous(), "\n                        f"{_target_shape_literal(reduction_output_name)})"\n                    )\n                continue\n            if op_type == "MEAN":\n                direct_expr = _direct_mean_reduction_expr(\n                    input_expr=_tensor_expr(str(op.inputs[0])),\n                    axes=axis_values,\n                    input_rank=len(list(model_ir.tensors[str(op.inputs[0])].shape)),\n                    keepdims=keepdims,\n                )\n                if direct_expr is not None:\n                    reduced_expr = _emit_maybe_aligned_expr(\n                        output_name=outputs[0],\n                        expr=direct_expr,\n                        inferred_shape=_infer_reduction_shape(_tensor_shape_list(str(op.inputs[0])), axis_values, keepdims=keepdims),\n                    )\n                    forward_lines.append(f"{reduction_raw_output_var} = {reduced_expr}")\n                    if reduction_raw_output_var != output_vars[0] and not _can_omit_materialized_channel_last_alias(reduction_output_name):\n                        runtime_imports.add("_align_tensor_to_target_shape")\n                        forward_lines.append(\n                            f"{output_vars[0]} = _align_tensor_to_target_shape("\n                            f"{reduction_raw_output_var}.permute(0, 2, 3, 1).contiguous(), "\n                            f"{_target_shape_literal(reduction_output_name)})"\n                        )\n                    continue\n            reduced_expr = _emit_maybe_aligned_expr(\n                output_name=outputs[0],\n                expr=f"{reducer_map[op_type]}({_tensor_expr(str(op.inputs[0]))}, {axis_expr}, {keepdims})",\n                inferred_shape=_infer_reduction_shape(_tensor_shape_list(str(op.inputs[0])), axis_values, keepdims=keepdims),\n            )\n            forward_lines.append(f"{reduction_raw_output_var} = {reduced_expr}")\n            if reduction_raw_output_var != output_vars[0] and not _can_omit_materialized_channel_last_alias(reduction_output_name):\n                runtime_imports.add("_align_tensor_to_target_shape")\n                forward_lines.append(\n                    f"{output_vars[0]} = _align_tensor_to_target_shape("\n                    f"{reduction_raw_output_var}.permute(0, 2, 3, 1).contiguous(), "\n                    f"{_target_shape_literal(reduction_output_name)})"\n                )\n            continue\n        if op_type == "PAD":\n            runtime_imports.add("_align_tensor_to_target_shape")\n            pad_input_expr = f"_align_tensor_to_target_shape({_tensor_expr(str(op.inputs[0]))}, {_target_shape_literal(str(op.inputs[0]))})"\n            pad_literal_expr = _pad_literal_expr(str(op.inputs[1]))\n            if pad_literal_expr is not None:\n                forward_lines.append(\n                    f"{output_vars[0]} = F.pad({pad_input_expr}, {pad_literal_expr}, mode=\'constant\', value=0.0)"\n                )\n            else:\n                runtime_imports.add("_apply_pad_nd")\n                forward_lines.append(\n                    f"{output_vars[0]} = _apply_pad_nd({pad_input_expr}, {_tensor_expr(str(op.inputs[1]))}, mode=\'constant\', value=0.0)"\n                )\n            continue\n        if op_type == "PADV2":\n            runtime_imports.add("_align_tensor_to_target_shape")\n            pad_input_expr = f"_align_tensor_to_target_shape({_tensor_expr(str(op.inputs[0]))}, {_target_shape_literal(str(op.inputs[0]))})"\n            value_expr = _scalar_literal_expr(str(op.inputs[2])) or f"float({_tensor_expr(str(op.inputs[2]))}.reshape(-1)[0].item())"\n            pad_literal_expr = _pad_literal_expr(str(op.inputs[1]))\n            if pad_literal_expr is not None:\n                forward_lines.append(\n                    f"{output_vars[0]} = F.pad({pad_input_expr}, {pad_literal_expr}, mode=\'constant\', value={value_expr})"\n                )\n            else:\n                runtime_imports.add("_apply_pad_nd")\n                forward_lines.append(\n                    f"{output_vars[0]} = _apply_pad_nd({pad_input_expr}, {_tensor_expr(str(op.inputs[1]))}, mode=\'constant\', value={value_expr})"\n                )\n            continue\n        if op_type == "MIRROR_PAD":\n            static_mirror_pad_expr = _static_mirror_pad_expr(\n                input_tensor_name=str(op.inputs[0]),\n                pads_tensor_name=str(op.inputs[1]),\n            )\n            if static_mirror_pad_expr is not None:\n                forward_lines.append(f"{output_vars[0]} = {static_mirror_pad_expr}")\n            else:\n                runtime_imports.add("_apply_pad_nd")\n                runtime_imports.add("_align_tensor_to_target_shape")\n                pad_input_expr = f"_align_tensor_to_target_shape({_tensor_expr(str(op.inputs[0]))}, {_target_shape_literal(str(op.inputs[0]))})"\n                forward_lines.append(\n                    f"{output_vars[0]} = _apply_pad_nd({pad_input_expr}, {_tensor_expr(str(op.inputs[1]))}, mode=\'reflect\')"\n                )\n            continue\n        if op_type in {"WHERE", "SELECT", "SELECT_V2"}:\n            if len(op.inputs) == 1:\n                forward_lines.append(\n                    f"{output_vars[0]} = torch.nonzero({_tensor_expr(str(op.inputs[0]))})"\n                )\n            else:\n                forward_lines.append(\n                    f"{output_vars[0]} = torch.where({_tensor_expr(str(op.inputs[0]))}, {_tensor_expr(str(op.inputs[1]))}, {_tensor_expr(str(op.inputs[2]))})"\n                )\n            continue\n        if op_type == "TILE":\n            runtime_imports.add("_apply_tile")\n            forward_lines.append(\n                f"{output_vars[0]} = _apply_tile({_tensor_expr(str(op.inputs[0]))}, {_tensor_expr(str(op.inputs[1]))})"\n            )\n            continue\n        if op_type == "BATCH_MATMUL":\n            x_expr = _tensor_expr(str(op.inputs[0]))\n            y_expr = _tensor_expr(str(op.inputs[1]))\n            adj_x = bool(op.options.get("adjX", False))\n            adj_y = bool(op.options.get("adjY", False))\n            effective_adj_x = bool(adj_x)\n            effective_adj_y = bool(adj_y)\n            x_shape = _tensor_shape_list(str(op.inputs[0]))\n            y_shape = _tensor_shape_list(str(op.inputs[1]))\n            output_shape = _tensor_shape_list(outputs[0])\n            lhs_producer_idx = producer_index.get(str(op.inputs[0]), None)\n            if lhs_producer_idx is not None:\n                lhs_producer = model_ir.operators[int(lhs_producer_idx)]\n                if str(lhs_producer.op_type) == "RESHAPE" and len(lhs_producer.inputs) >= 1 and len(lhs_producer.outputs) == 1:\n                    lhs_feature_last_target = _reshape_prefers_feature_last_for_adjx_batch_matmul(\n                        str(lhs_producer.inputs[0]),\n                        str(lhs_producer.outputs[0]),\n                    )\n                    if lhs_feature_last_target is not None:\n                        x_shape = [int(v) for v in list(lhs_feature_last_target[1])]\n            inferred_shape = _infer_batch_matmul_shape(\n                x_shape,\n                y_shape,\n                adj_x=effective_adj_x,\n                adj_y=effective_adj_y,\n            )\n            if output_shape is not None and not _shape_lists_equal(inferred_shape, output_shape):\n                best_flag_choice: Optional[Tuple[int, bool, bool, List[int]]] = None\n                for candidate_adj_x, candidate_adj_y in [\n                    (effective_adj_x, effective_adj_y),\n                    (effective_adj_x, not effective_adj_y),\n                    (not effective_adj_x, effective_adj_y),\n                    (not effective_adj_x, not effective_adj_y),\n                ]:\n                    candidate_shape = _infer_batch_matmul_shape(\n                        x_shape,\n                        y_shape,\n                        adj_x=candidate_adj_x,\n                        adj_y=candidate_adj_y,\n                    )\n                    if not _shape_lists_equal(candidate_shape, output_shape):\n                        continue\n                    score = int(candidate_adj_x != adj_x) + int(candidate_adj_y != adj_y)\n                    choice = (int(score), bool(candidate_adj_x), bool(candidate_adj_y), [int(v) for v in list(candidate_shape or [])])\n                    if best_flag_choice is None or choice < best_flag_choice:\n                        best_flag_choice = choice\n                if best_flag_choice is not None:\n                    _, effective_adj_x, effective_adj_y, inferred_shape = best_flag_choice\n            emitted_x_expr = x_expr\n            emitted_y_expr = y_expr\n            runtime_adj_x = bool(effective_adj_x)\n            runtime_adj_y = bool(effective_adj_y)\n            if runtime_adj_x:\n                transposed_x_expr = _transposed_constant_expr_for_tensor_name(str(op.inputs[0]))\n                if transposed_x_expr is not None:\n                    emitted_x_expr = str(transposed_x_expr)\n                    runtime_adj_x = False\n            if runtime_adj_y:\n                transposed_y_expr = _transposed_constant_expr_for_tensor_name(str(op.inputs[1]))\n                if transposed_y_expr is not None:\n                    emitted_y_expr = str(transposed_y_expr)\n                    runtime_adj_y = False\n            forward_lines.append(f"_tmp_x_{op_index} = {emitted_x_expr}")\n            forward_lines.append(f"_tmp_y_{op_index} = {emitted_y_expr}")\n            if runtime_adj_x:\n                forward_lines.append(f"_tmp_x_{op_index} = _tmp_x_{op_index}.transpose(-1, -2)")\n            if runtime_adj_y:\n                forward_lines.append(f"_tmp_y_{op_index} = _tmp_y_{op_index}.transpose(-1, -2)")\n            if (\n                not effective_adj_x\n                and not effective_adj_y\n                and x_shape is not None\n                and y_shape is not None\n                and len(x_shape) >= 3\n                and len(y_shape) == 2\n            ):\n                transposed_x_shape = list(x_shape[:-2]) + [int(x_shape[-1]), int(x_shape[-2])]\n                inferred_shape_with_x_transpose = _infer_batch_matmul_shape(\n                    transposed_x_shape,\n                    y_shape,\n                    adj_x=False,\n                    adj_y=False,\n                )\n                input_tensor = model_ir.tensors.get(str(op.inputs[0]), None)\n                input_layout = normalize_logical_layout(input_tensor.logical_layout) if input_tensor is not None else LOGICAL_LAYOUT_UNKNOWN\n                if (\n                    inferred_shape is None\n                    and inferred_shape_with_x_transpose is not None\n                    and is_channel_first_logical_layout(input_layout)\n                ):\n                    forward_lines.append(f"_tmp_x_{op_index} = _tmp_x_{op_index}.transpose(-1, -2)")\n                    inferred_shape = inferred_shape_with_x_transpose\n            runtime_imports.add("_align_tensor_to_target_shape")\n            if output_target_shape != "None":\n                forward_lines.append(\n                    f"{output_vars[0]} = _align_tensor_to_target_shape(torch.matmul(_tmp_x_{op_index}, _tmp_y_{op_index}), {output_target_shape})"\n                )\n            else:\n                forward_lines.append(\n                    f"{output_vars[0]} = {_emit_maybe_aligned_expr(output_name=outputs[0], expr=f\'torch.matmul(_tmp_x_{op_index}, _tmp_y_{op_index})\', inferred_shape=inferred_shape)}"\n                )\n            continue\n        if op_type == "NON_MAX_SUPPRESSION_V4":\n            runtime_imports.add("_apply_non_max_suppression_v4")\n            nms_method_name = f"_run_nms_{len(nms_method_specs)}"\n            max_output_literal = _int_scalar_literal_expr(str(op.inputs[2]))\n            max_output_shape_expr = None if max_output_literal is not None else _reconstruct_shape_scalar_expr(str(op.inputs[2]))\n            if max_output_shape_expr is not None:\n                max_output_shape_expr = max_output_shape_expr.replace(_tensor_expr(str(op.inputs[0])), "boxes")\n                max_output_shape_expr = max_output_shape_expr.replace(_tensor_expr(str(op.inputs[1])), "scores")\n            iou_threshold_literal = _scalar_literal_expr(str(op.inputs[3]))\n            score_threshold_literal = _scalar_literal_expr(str(op.inputs[4]))\n            nms_method_specs.append(\n                {\n                    "name": nms_method_name,\n                    "max_output_expr": max_output_literal or max_output_shape_expr or "max_output_size",\n                    "iou_threshold_expr": iou_threshold_literal or "iou_threshold",\n                    "score_threshold_expr": score_threshold_literal or "score_threshold",\n                    "max_output_arg_expr": None if (max_output_literal is not None or max_output_shape_expr is not None) else _tensor_expr(str(op.inputs[2])),\n                    "iou_threshold_arg_expr": None if iou_threshold_literal is not None else _tensor_expr(str(op.inputs[3])),\n                    "score_threshold_arg_expr": None if score_threshold_literal is not None else _tensor_expr(str(op.inputs[4])),\n                }\n            )\n            call_args = [\n                _tensor_expr(str(op.inputs[0])),\n                _tensor_expr(str(op.inputs[1])),\n            ]\n            if max_output_literal is None and max_output_shape_expr is None:\n                call_args.append(_tensor_expr(str(op.inputs[2])))\n            if iou_threshold_literal is None:\n                call_args.append(_tensor_expr(str(op.inputs[3])))\n            if score_threshold_literal is None:\n                call_args.append(_tensor_expr(str(op.inputs[4])))\n            forward_lines.append(\n                f"{\', \'.join(output_vars)} = self.{nms_method_name}({\', \'.join(call_args)})"\n            )\n            runtime_shape_uncertain_tensors.update(outputs)\n            if len(outputs) >= 2 and len(output_vars) >= 2 and len(consumer_index.get(outputs[0], [])) > 0:\n                runtime_imports.add("_crop_nms_selected_indices")\n                cropped_indices_var = f"_nms_selected_indices_valid_{op_index}"\n                forward_lines.append(\n                    f"{cropped_indices_var} = _crop_nms_selected_indices({output_vars[0]}, {output_vars[1]})"\n                )\n                tensor_expr_aliases[outputs[0]] = cropped_indices_var\n            continue\n        raise ModelIRPyTorchExportError(\n            "Native PyTorch-like model.py codegen hit an unimplemented op emitter. "\n            f"op_type={op_type}"\n        )\n\n    has_conv_blocks = len(fused_module_specs) > 0\n    conv_block_helper_source = (\n        "class _Conv2dBlock(torch.nn.Module):\\n"\n        "    def __init__(\\n"\n        "        self,\\n"\n        "        *,\\n"\n        "        in_channels: int,\\n"\n        "        out_channels: int,\\n"\n        "        kernel_size: tuple[int, int],\\n"\n        "        stride: tuple[int, int],\\n"\n        "        padding: tuple[int, int],\\n"\n        "        dilation: tuple[int, int],\\n"\n        "        groups: int,\\n"\n        "        bias: bool,\\n"\n        "        pad: Optional[list[int]] = None,\\n"\n        "        activation: str = \'none\',\\n"\n        "        negative_slope: float = 0.2,\\n"\n        "        pad_mode: str = \'constant\',\\n"\n        "        pad_value: float = 0.0,\\n"\n        "    ) -> None:\\n"\n        "        super().__init__()\\n"\n        "        self.conv = torch.nn.Conv2d(\\n"\n        "            in_channels=in_channels,\\n"\n        "            out_channels=out_channels,\\n"\n        "            kernel_size=kernel_size,\\n"\n        "            stride=stride,\\n"\n        "            padding=padding,\\n"\n        "            dilation=dilation,\\n"\n        "            groups=groups,\\n"\n        "            bias=bias,\\n"\n        "        )\\n"\n        "        self.pad = pad\\n"\n        "        self.activation = str(activation)\\n"\n        "        self.negative_slope = float(negative_slope)\\n"\n        "        self.pad_mode = str(pad_mode)\\n"\n        "        self.pad_value = float(pad_value)\\n\\n"\n        "    def forward(self, x: torch.Tensor) -> torch.Tensor:\\n"\n        "        input_was_channel_last = False\\n"\n        "        if x.ndim == 4 and int(x.shape[1]) != int(self.conv.in_channels) and int(x.shape[-1]) == int(self.conv.in_channels):\\n"\n        "            x = x.permute(0, 3, 1, 2).contiguous()\\n"\n        "            input_was_channel_last = True\\n"\n        "        if self.pad is not None:\\n"\n        "            x = F.pad(x, self.pad, mode=self.pad_mode, value=self.pad_value)\\n"\n        "        x = self.conv(x)\\n"\n        "        if input_was_channel_last:\\n"\n        "            x = x.permute(0, 2, 3, 1).contiguous()\\n"\n        "        if self.activation == \'leaky_relu\':\\n"\n        "            return F.leaky_relu(x, negative_slope=self.negative_slope)\\n"\n        "        if self.activation == \'relu\':\\n"\n        "            return torch.relu(x)\\n"\n        "        if self.activation == \'relu6\':\\n"\n        "            return torch.clamp(x, min=0.0, max=6.0)\\n"\n        "        if self.activation == \'relu_n1_to_1\':\\n"\n        "            return torch.clamp(x, min=-1.0, max=1.0)\\n"\n        "        if self.activation == \'relu_0_to_1\':\\n"\n        "            return torch.clamp(x, min=0.0, max=1.0)\\n"\n        "        if self.activation == \'tanh\':\\n"\n        "            return torch.tanh(x)\\n"\n        "        if self.activation == \'sigmoid\':\\n"\n        "            return torch.sigmoid(x)\\n"\n        "        if self.activation == \'silu\':\\n"\n        "            return torch.mul(x, torch.sigmoid(x))\\n"\n        "        return x\\n\\n"\n    ) if has_conv_blocks else ""\n    has_sequence_lstm_blocks = any(\n        str(op.op_type) in {"UNIDIRECTIONAL_SEQUENCE_LSTM", "BIDIRECTIONAL_SEQUENCE_LSTM"}\n        for op in model_ir.operators\n    )\n    has_sequence_rnn_blocks = any(\n        str(op.op_type) == "UNIDIRECTIONAL_SEQUENCE_RNN"\n        for op in model_ir.operators\n    )\n    sequence_rnn_helper_source = (\n        "class _SequenceRNNBlock(torch.nn.Module):\\n"\n        "    def __init__(self, *, input_size: int, hidden_size: int, activation: str) -> None:\\n"\n        "        super().__init__()\\n"\n        "        activation_key = str(activation).strip().lower()\\n"\n        "        if activation_key not in {\'tanh\', \'relu\'}:\\n"\n        "            raise ValueError(f\'Unsupported RNN activation for native codegen: {activation}\')\\n"\n        "        self.rnn = torch.nn.RNN(\\n"\n        "            input_size=int(input_size),\\n"\n        "            hidden_size=int(hidden_size),\\n"\n        "            num_layers=1,\\n"\n        "            nonlinearity=activation_key,\\n"\n        "            bias=True,\\n"\n        "            batch_first=False,\\n"\n        "            bidirectional=False,\\n"\n        "        )\\n\\n"\n        "    def forward(\\n"\n        "        self,\\n"\n        "        x: torch.Tensor,\\n"\n        "        h0: Optional[torch.Tensor] = None,\\n"\n        "    ) -> torch.Tensor:\\n"\n        "        state: Optional[torch.Tensor] = None\\n"\n        "        if h0 is not None:\\n"\n        "            state = torch.unsqueeze(h0, dim=0)\\n"\n        "        y, _ = self.rnn(x, state) if state is not None else self.rnn(x)\\n"\n        "        return y\\n\\n"\n    ) if has_sequence_rnn_blocks else ""\n    sequence_lstm_helper_source = (\n        "class _SequenceLSTMBlock(torch.nn.Module):\\n"\n        "    def __init__(self, *, input_size: int, hidden_size: int, sequence_length: Optional[int], bidirectional: bool, merge_outputs: bool) -> None:\\n"\n        "        super().__init__()\\n"\n        "        self.hidden_size = int(hidden_size)\\n"\n        "        self.sequence_length = int(sequence_length) if sequence_length is not None and int(sequence_length) > 0 else None\\n"\n        "        self.bidirectional = bool(bidirectional)\\n"\n        "        self.merge_outputs = bool(merge_outputs)\\n"\n        "        self.lstm = torch.nn.LSTM(\\n"\n        "            input_size=int(input_size),\\n"\n        "            hidden_size=int(hidden_size),\\n"\n        "            num_layers=1,\\n"\n        "            bias=True,\\n"\n        "            batch_first=False,\\n"\n        "            bidirectional=bool(bidirectional),\\n"\n        "        )\\n\\n"\n        "    def forward(\\n"\n        "        self,\\n"\n        "        x: torch.Tensor,\\n"\n        "        fw_h0: Optional[torch.Tensor] = None,\\n"\n        "        fw_c0: Optional[torch.Tensor] = None,\\n"\n        "        bw_h0: Optional[torch.Tensor] = None,\\n"\n        "        bw_c0: Optional[torch.Tensor] = None,\\n"\n        "    ) -> torch.Tensor:\\n"\n        "        if self.sequence_length is not None:\\n"\n        "            x = torch.reshape(x, (self.sequence_length, x.shape[1], self.lstm.input_size))\\n"\n        "        else:\\n"\n        "            x = torch.reshape(x, (x.shape[0], x.shape[1], self.lstm.input_size))\\n"\n        "        if not self.bidirectional:\\n"\n        "            state: Optional[Tuple[torch.Tensor, torch.Tensor]] = None\\n"\n        "            if fw_h0 is not None and fw_c0 is not None:\\n"\n        "                state = (torch.unsqueeze(fw_h0, dim=0), torch.unsqueeze(fw_c0, dim=0))\\n"\n        "            y, _ = self.lstm(x, state) if state is not None else self.lstm(x)\\n"\n        "            return y\\n"\n        "        state: Optional[Tuple[torch.Tensor, torch.Tensor]] = None\\n"\n        "        if fw_h0 is not None or fw_c0 is not None or bw_h0 is not None or bw_c0 is not None:\\n"\n        "            batch_size = x.shape[1]\\n"\n        "            zeros = torch.zeros((batch_size, self.hidden_size), dtype=x.dtype, device=x.device)\\n"\n        "            fw_h = fw_h0 if fw_h0 is not None else zeros\\n"\n        "            fw_c = fw_c0 if fw_c0 is not None else zeros\\n"\n        "            bw_h = bw_h0 if bw_h0 is not None else zeros\\n"\n        "            bw_c = bw_c0 if bw_c0 is not None else zeros\\n"\n        "            fw_h = torch.reshape(fw_h, (batch_size, self.hidden_size))\\n"\n        "            fw_c = torch.reshape(fw_c, (batch_size, self.hidden_size))\\n"\n        "            bw_h = torch.reshape(bw_h, (batch_size, self.hidden_size))\\n"\n        "            bw_c = torch.reshape(bw_c, (batch_size, self.hidden_size))\\n"\n        "            state = (\\n"\n        "                torch.stack([fw_h, bw_h], dim=0),\\n"\n        "                torch.stack([fw_c, bw_c], dim=0),\\n"\n        "            )\\n"\n        "        y, _ = self.lstm(x, state) if state is not None else self.lstm(x)\\n"\n        "        if self.merge_outputs:\\n"\n        "            return y\\n"\n        "        return torch.stack([y[:, :, : self.hidden_size], y[:, :, self.hidden_size :]], dim=2)\\n\\n"\n    ) if has_sequence_lstm_blocks else ""\n    helper_source = (\n        "def _normalize_tensor_name(name: str) -> str:\\n"\n        "    normalized = str(name).split(\\":\\")[0]\\n"\n        "    if normalized.startswith(\\"serving_default_\\"):\\n"\n        "        normalized = normalized[len(\\"serving_default_\\") :]\\n"\n        "    return normalized\\n\\n"\n        "_TORCH_DTYPE_BY_TFLITE_DTYPE: Dict[str, torch.dtype] = {\\n"\n        "    \'BOOL\': torch.bool,\\n"\n        "    \'INT8\': torch.int8,\\n"\n        "    \'INT16\': torch.int16,\\n"\n        "    \'INT32\': torch.int32,\\n"\n        "    \'INT64\': torch.int64,\\n"\n        "    \'UINT8\': torch.uint8,\\n"\n        "    \'FLOAT16\': torch.float16,\\n"\n        "    \'FLOAT32\': torch.float32,\\n"\n        "    \'FLOAT64\': torch.float64,\\n"\n        "}\\n\\n"\n        "_ONNX2TF_DISABLE_SYMBOLIC_SHAPE_TENSORS = False\\n\\n"\n        "def _torch_dtype(dtype_name: str) -> torch.dtype:\\n"\n        "    key = str(dtype_name).upper()\\n"\n        "    if key not in _TORCH_DTYPE_BY_TFLITE_DTYPE:\\n"\n        "        raise RuntimeError(f\'Unsupported dtype for PyTorch runtime: {dtype_name}\')\\n"\n        "    return _TORCH_DTYPE_BY_TFLITE_DTYPE[key]\\n\\n"\n        "def _module_device(module: Any) -> torch.device:\\n"\n        "    if torch.jit.is_scripting():\\n"\n        "        return torch.device(\'cpu\')\\n"\n        "    for parameter in module.parameters():\\n"\n        "        return parameter.device\\n"\n        "    for buffer in module.buffers():\\n"\n        "        return buffer.device\\n"\n        "    return torch.device(\'cpu\')\\n\\n"\n        "def _default_tensor_storage_name(tensor_name: str) -> str:\\n"\n        "    base_name = re.sub(r\'[^0-9A-Za-z_]\', \'_\', str(tensor_name)).strip(\'_\')\\n"\n        "    if base_name == \'\':\\n"\n        "        base_name = \'tensor\'\\n"\n        "    if base_name[0].isdigit():\\n"\n        "        base_name = f\'tensor_{base_name}\'\\n"\n        "    return base_name\\n\\n"\n        "def _resolve_named_input_value(kwargs: Dict[str, Any], expected_name: str) -> Any:\\n"\n        "    if str(expected_name) in kwargs:\\n"\n        "        return kwargs[str(expected_name)]\\n"\n        "    normalized_expected_name = _normalize_tensor_name(str(expected_name))\\n"\n        "    canonical_expected_name = re.sub(r\'[^0-9A-Za-z]+\', \'_\', str(expected_name)).strip(\'_\').lower()\\n"\n        "    for candidate_name, candidate_value in kwargs.items():\\n"\n        "        normalized_candidate = _normalize_tensor_name(str(candidate_name))\\n"\n        "        canonical_candidate = re.sub(r\'[^0-9A-Za-z]+\', \'_\', str(candidate_name)).strip(\'_\').lower()\\n"\n        "        if (\\n"\n        "            normalized_candidate == normalized_expected_name\\n"\n        "            or canonical_candidate == canonical_expected_name\\n"\n        "            or normalized_candidate.endswith(normalized_expected_name)\\n"\n        "            or normalized_expected_name.endswith(normalized_candidate)\\n"\n        "            or canonical_candidate.endswith(canonical_expected_name)\\n"\n        "            or canonical_expected_name.endswith(canonical_candidate)\\n"\n        "        ):\\n"\n        "            return candidate_value\\n"\n        "    raise KeyError(str(expected_name))\\n\\n"\n        "def resolve_named_input_value(kwargs: Dict[str, Any], expected_name: str) -> Any:\\n"\n        "    return _resolve_named_input_value(kwargs, expected_name)\\n\\n"\n        "def _perm_cl_to_cf(rank: int) -> Optional[List[int]]:\\n"\n        "    if rank == 3:\\n"\n        "        return [0, 2, 1]\\n"\n        "    if rank == 4:\\n"\n        "        return [0, 3, 1, 2]\\n"\n        "    if rank == 5:\\n"\n        "        return [0, 4, 1, 2, 3]\\n"\n        "    return None\\n\\n"\n        "def _perm_cf_to_cl(rank: int) -> Optional[List[int]]:\\n"\n        "    if rank == 3:\\n"\n        "        return [0, 2, 1]\\n"\n        "    if rank == 4:\\n"\n        "        return [0, 2, 3, 1]\\n"\n        "    if rank == 5:\\n"\n        "        return [0, 2, 3, 4, 1]\\n"\n        "    return None\\n\\n"\n        "def _permute_shape(values: Sequence[int], perm: Sequence[int]) -> List[int]:\\n"\n        "    items = [int(v) for v in list(values)]\\n"\n        "    return [int(items[idx]) for idx in perm]\\n\\n"\n        "def _torch_permute(value: torch.Tensor, perm: List[int]) -> torch.Tensor:\\n"\n        "    if len(perm) == 0:\\n"\n        "        return value\\n"\n        "    return torch.permute(value, perm).contiguous()\\n\\n"\n        "def _can_broadcast_shapes(lhs: List[int], rhs: List[int]) -> bool:\\n"\n        "    lhs_index = len(lhs) - 1\\n"\n        "    rhs_index = len(rhs) - 1\\n"\n        "    while lhs_index >= 0 and rhs_index >= 0:\\n"\n        "        lhs_dim = int(lhs[lhs_index])\\n"\n        "        rhs_dim = int(rhs[rhs_index])\\n"\n        "        if lhs_dim != rhs_dim and lhs_dim != 1 and rhs_dim != 1:\\n"\n        "            return False\\n"\n        "        lhs_index -= 1\\n"\n        "        rhs_index -= 1\\n"\n        "    return True\\n\\n"\n        "def _broadcast_shape(lhs: Sequence[int], rhs: Sequence[int]) -> Optional[List[int]]:\\n"\n        "    lhs_items = [int(v) for v in list(lhs)]\\n"\n        "    rhs_items = [int(v) for v in list(rhs)]\\n"\n        "    result_rev = torch.jit.annotate(List[int], [])\\n"\n        "    lhs_index = len(lhs_items) - 1\\n"\n        "    rhs_index = len(rhs_items) - 1\\n"\n        "    while lhs_index >= 0 or rhs_index >= 0:\\n"\n        "        lhs_dim = int(lhs_items[lhs_index]) if lhs_index >= 0 else 1\\n"\n        "        rhs_dim = int(rhs_items[rhs_index]) if rhs_index >= 0 else 1\\n"\n        "        if lhs_dim != rhs_dim and lhs_dim != 1 and rhs_dim != 1:\\n"\n        "            return None\\n"\n        "        result_rev.append(int(lhs_dim if rhs_dim == 1 else rhs_dim))\\n"\n        "        lhs_index -= 1\\n"\n        "        rhs_index -= 1\\n"\n        "    result = torch.jit.annotate(List[int], [])\\n"\n        "    reverse_index = len(result_rev) - 1\\n"\n        "    while reverse_index >= 0:\\n"\n        "        result.append(int(result_rev[reverse_index]))\\n"\n        "        reverse_index -= 1\\n"\n        "    return result\\n\\n"\n        "def _align_tensor_to_target_shape(value: torch.Tensor, target_shape: Optional[Sequence[int]]) -> torch.Tensor:\\n"\n        "    if target_shape is None:\\n"\n        "        return value\\n"\n        "    has_actual_shape, actual_shape = _tensor_static_shape_list(value)\\n"\n        "    if not has_actual_shape:\\n"\n        "        return value\\n"\n        "    target = [int(v) for v in list(target_shape)]\\n"\n        "    if actual_shape == target:\\n"\n        "        return value\\n"\n        "    perm = _perm_cl_to_cf(value.ndim)\\n"\n        "    if perm is not None and _permute_shape(actual_shape, perm) == target:\\n"\n        "        return _torch_permute(value, perm)\\n"\n        "    perm_inv = _perm_cf_to_cl(value.ndim)\\n"\n        "    if perm_inv is not None and _permute_shape(actual_shape, perm_inv) == target:\\n"\n        "        return _torch_permute(value, perm_inv)\\n"\n        "    if len(actual_shape) == len(target):\\n"\n        "        if all(int(dim) == 1 for dim in list(target)) and any(int(actual_dim) > 1 for actual_dim in list(actual_shape)):\\n"\n        "            return value\\n"\n        "        can_narrow = True\\n"\n        "        has_mismatch = False\\n"\n        "        for dim_idx, (actual_dim, target_dim) in enumerate(zip(actual_shape, target)):\\n"\n        "            if int(target_dim) <= 0 or int(actual_dim) < int(target_dim):\\n"\n        "                can_narrow = False\\n"\n        "                break\\n"\n        "            if int(actual_dim) != int(target_dim):\\n"\n        "                has_mismatch = True\\n"\n        "                if int(dim_idx) == 0:\\n"\n        "                    can_narrow = False\\n"\n        "                    break\\n"\n        "        if can_narrow and has_mismatch:\\n"\n        "            narrowed = value\\n"\n        "            for dim_idx, (actual_dim, target_dim) in enumerate(zip(actual_shape, target)):\\n"\n        "                if int(actual_dim) > int(target_dim):\\n"\n        "                    narrowed = torch.narrow(narrowed, int(dim_idx), 0, int(target_dim))\\n"\n        "            return narrowed\\n"\n        "    return value\\n\\n"\n        "@torch.jit.ignore\\n"\n        "def _align_scatter_nd_updates_eager(updates: torch.Tensor, expected_shape: Sequence[int]) -> torch.Tensor:\\n"\n        "    has_actual_shape, actual_shape = _tensor_static_shape_list(updates)\\n"\n        "    if not has_actual_shape:\\n"\n        "        return updates\\n"\n        "    expected = []\\n"\n        "    for dim in list(expected_shape):\\n"\n        "        if not isinstance(dim, int):\\n"\n        "            return updates\\n"\n        "        expected.append(int(dim))\\n"\n        "    try:\\n"\n        "        if list(torch.broadcast_shapes(tuple(actual_shape), tuple(expected))) == expected:\\n"\n        "            return updates\\n"\n        "    except Exception:\\n"\n        "        pass\\n"\n        "    perm = _perm_cf_to_cl(updates.ndim)\\n"\n        "    if perm is not None and _permute_shape(actual_shape, perm) == expected:\\n"\n        "        return _torch_permute(updates, perm)\\n"\n        "    perm = _perm_cl_to_cf(updates.ndim)\\n"\n        "    if perm is not None and _permute_shape(actual_shape, perm) == expected:\\n"\n        "        return _torch_permute(updates, perm)\\n"\n        "    if updates.ndim <= 5:\\n"\n        "        import itertools\\n"\n        "        for generic_perm in itertools.permutations(range(updates.ndim)):\\n"\n        "            if list(generic_perm) == list(range(updates.ndim)):\\n"\n        "                continue\\n"\n        "            generic_perm_list = [int(v) for v in list(generic_perm)]\\n"\n        "            if _permute_shape(actual_shape, generic_perm_list) == expected:\\n"\n        "                return _torch_permute(updates, generic_perm_list)\\n"\n        "    return updates\\n\\n"\n        "def _align_scatter_nd_updates(updates: torch.Tensor, expected_shape: Sequence[int]) -> torch.Tensor:\\n"\n        "    has_actual_shape, actual_shape = _tensor_static_shape_list(updates)\\n"\n        "    expected = torch.jit.annotate(List[int], [])\\n"\n        "    for dim in list(expected_shape):\\n"\n        "        if not isinstance(dim, int):\\n"\n        "            return updates\\n"\n        "        expected.append(int(dim))\\n"\n        "    if not has_actual_shape:\\n"\n        "        return updates\\n"\n        "    if actual_shape == expected:\\n"\n        "        return updates\\n"\n        "    broadcast_shape = _broadcast_shape(actual_shape, expected)\\n"\n        "    if broadcast_shape is not None and [int(v) for v in list(broadcast_shape)] == expected:\\n"\n        "        return updates\\n"\n        "    perm = _perm_cf_to_cl(updates.ndim)\\n"\n        "    if perm is not None and _permute_shape(actual_shape, perm) == expected:\\n"\n        "        return _torch_permute(updates, perm)\\n"\n        "    perm = _perm_cl_to_cf(updates.ndim)\\n"\n        "    if perm is not None and _permute_shape(actual_shape, perm) == expected:\\n"\n        "        return _torch_permute(updates, perm)\\n"\n        "    if torch.jit.is_scripting():\\n"\n        "        return updates\\n"\n        "    return _align_scatter_nd_updates_eager(updates, expected)\\n\\n"\n        "def _matches_target_except_axis(actual_shape: Sequence[int], target_shape: Sequence[int], axis: int) -> bool:\\n"\n        "    if len(list(actual_shape)) != len(list(target_shape)):\\n"\n        "        return False\\n"\n        "    for idx, (actual_dim, target_dim) in enumerate(zip(actual_shape, target_shape)):\\n"\n        "        if int(idx) == int(axis):\\n"\n        "            continue\\n"\n        "        if int(target_dim) <= 1:\\n"\n        "            continue\\n"\n        "        if int(actual_dim) != int(target_dim):\\n"\n        "            return False\\n"\n        "    return True\\n\\n"\n        "def _optional_static_shape_list(shape: Optional[Sequence[int]]) -> Tuple[bool, List[int]]:\\n"\n        "    if shape is None:\\n"\n        "        return False, []\\n"\n        "    values = torch.jit.annotate(List[int], [])\\n"\n        "    for dim in list(shape):\\n"\n        "        if not isinstance(dim, int):\\n"\n        "            return False, []\\n"\n        "        values.append(int(dim))\\n"\n        "    return True, values\\n\\n"\n        "def _shape_matches_target_relaxed(actual_shape: Sequence[int], target_shape: Sequence[int]) -> bool:\\n"\n        "    if len(list(actual_shape)) != len(list(target_shape)):\\n"\n        "        return False\\n"\n        "    for actual_dim, target_dim in zip(actual_shape, target_shape):\\n"\n        "        if int(target_dim) <= 0:\\n"\n        "            continue\\n"\n        "        if int(actual_dim) != int(target_dim):\\n"\n        "            return False\\n"\n        "    return True\\n\\n"\n        "def _target_shape_has_concrete_dim(target_shape: Optional[Sequence[int]]) -> bool:\\n"\n        "    if target_shape is None:\\n"\n        "        return False\\n"\n        "    for dim in list(target_shape):\\n"\n        "        if int(dim) > 0:\\n"\n        "            return True\\n"\n        "    return False\\n\\n"\n        "def _tensor_static_shape_list(value: torch.Tensor) -> Tuple[bool, List[int]]:\\n"\n        "    values = torch.jit.annotate(List[int], [])\\n"\n        "    if torch.jit.is_scripting():\\n"\n        "        for dim in list(value.shape):\\n"\n        "            values.append(int(dim))\\n"\n        "        return True, values\\n"\n        "    for dim in list(value.shape):\\n"\n        "        if not isinstance(dim, int):\\n"\n        "            return False, []\\n"\n        "        values.append(int(dim))\\n"\n        "    return True, values\\n\\n"\n        "@torch.jit.ignore\\n"\n        "def _tensor_shape_list_eager(value: torch.Tensor) -> List[int]:\\n"\n        "    return list(value.shape)\\n\\n"\n        "def _tensor_shape_list(value: torch.Tensor) -> List[int]:\\n"\n        "    if torch.jit.is_scripting():\\n"\n        "        result = torch.jit.annotate(List[int], [])\\n"\n        "        for dim in list(value.shape):\\n"\n        "            result.append(int(dim))\\n"\n        "        return result\\n"\n        "    return _tensor_shape_list_eager(value)\\n\\n"\n        "@torch.jit.ignore\\n"\n        "def _align_binary_inputs_eager(x: torch.Tensor, y: torch.Tensor, target_shape: Optional[Sequence[int]]) -> Tuple[torch.Tensor, torch.Tensor]:\\n"\n        "    target = [int(v) for v in list(target_shape)] if target_shape is not None else None\\n"\n        "    try:\\n"\n        "        broadcast_shape = list(torch.broadcast_shapes(tuple(int(v) for v in x.shape), tuple(int(v) for v in y.shape)))\\n"\n        "        if target is None or _shape_matches_target_relaxed(broadcast_shape, target):\\n"\n            "            return x, y\\n"\n        "    except Exception:\\n"\n            "        pass\\n"\n        "    perm = _perm_cl_to_cf(x.ndim)\\n"\n        "    if perm is None:\\n"\n        "        return x, y\\n"\n        "    x_shape = [int(v) for v in list(x.shape)]\\n"\n        "    y_shape = [int(v) for v in list(y.shape)]\\n"\n        "    if _permute_shape(y_shape, perm) == x_shape:\\n"\n        "        return _align_binary_inputs_to_anchor(x, _torch_permute(y, perm), target_shape)\\n"\n        "    if _permute_shape(x_shape, perm) == y_shape:\\n"\n        "        return _align_binary_inputs_to_anchor(_torch_permute(x, perm), y, target_shape)\\n"\n        "    if target is not None:\\n"\n        "        if _shape_matches_target_relaxed(_permute_shape(y_shape, perm), target):\\n"\n        "            return _align_binary_inputs_to_anchor(x, _torch_permute(y, perm), target_shape)\\n"\n        "        if _shape_matches_target_relaxed(_permute_shape(x_shape, perm), target):\\n"\n        "            return _align_binary_inputs_to_anchor(_torch_permute(x, perm), y, target_shape)\\n"\n        "        if x.ndim <= 5:\\n"\n        "            import itertools\\n"\n        "            for generic_perm in itertools.permutations(range(x.ndim)):\\n"\n        "                if list(generic_perm) == list(range(x.ndim)):\\n"\n        "                    continue\\n"\n        "                generic_perm_list = [int(v) for v in list(generic_perm)]\\n"\n        "                permuted_y_shape = _permute_shape(y_shape, generic_perm_list)\\n"\n        "                if permuted_y_shape is not None:\\n"\n        "                    try:\\n"\n        "                        broadcast_shape = list(torch.broadcast_shapes(tuple(permuted_y_shape), tuple(x_shape)))\\n"\n        "                        if _shape_matches_target_relaxed(broadcast_shape, target):\\n"\n        "                            return _align_binary_inputs_to_anchor(x, _torch_permute(y, generic_perm_list), target_shape)\\n"\n        "                    except Exception:\\n"\n        "                        pass\\n"\n        "                permuted_x_shape = _permute_shape(x_shape, generic_perm_list)\\n"\n        "                if permuted_x_shape is not None:\\n"\n        "                    try:\\n"\n        "                        broadcast_shape = list(torch.broadcast_shapes(tuple(permuted_x_shape), tuple(y_shape)))\\n"\n        "                        if _shape_matches_target_relaxed(broadcast_shape, target):\\n"\n        "                            return _align_binary_inputs_to_anchor(_torch_permute(x, generic_perm_list), y, target_shape)\\n"\n        "                    except Exception:\\n"\n        "                        pass\\n"\n        "    if x.ndim <= 5:\\n"\n        "        import itertools\\n"\n        "        for generic_perm in itertools.permutations(range(x.ndim)):\\n"\n        "            if list(generic_perm) == list(range(x.ndim)):\\n"\n        "                continue\\n"\n        "            generic_perm_list = [int(v) for v in list(generic_perm)]\\n"\n        "            permuted_y_shape = _permute_shape(y_shape, generic_perm_list)\\n"\n        "            if permuted_y_shape is not None:\\n"\n        "                try:\\n"\n        "                    torch.broadcast_shapes(tuple(permuted_y_shape), tuple(x_shape))\\n"\n        "                    return x, _torch_permute(y, generic_perm_list)\\n"\n        "                except Exception:\\n"\n        "                    pass\\n"\n        "                if target is not None:\\n"\n        "                    try:\\n"\n        "                        torch.broadcast_shapes(tuple(permuted_y_shape), tuple(target))\\n"\n        "                        return x, _torch_permute(y, generic_perm_list)\\n"\n        "                    except Exception:\\n"\n        "                        pass\\n"\n        "            permuted_x_shape = _permute_shape(x_shape, generic_perm_list)\\n"\n        "            if permuted_x_shape is not None:\\n"\n        "                try:\\n"\n        "                    torch.broadcast_shapes(tuple(permuted_x_shape), tuple(y_shape))\\n"\n        "                    return _torch_permute(x, generic_perm_list), y\\n"\n        "                except Exception:\\n"\n        "                    pass\\n"\n        "                if target is not None:\\n"\n        "                    try:\\n"\n        "                        torch.broadcast_shapes(tuple(permuted_x_shape), tuple(target))\\n"\n        "                        return _torch_permute(x, generic_perm_list), y\\n"\n        "                    except Exception:\\n"\n        "                        pass\\n"\n        "    return x, y\\n\\n"\n        "def _align_binary_inputs(x: torch.Tensor, y: torch.Tensor, target_shape: Optional[Sequence[int]]) -> Tuple[torch.Tensor, torch.Tensor]:\\n"\n        "    has_target, target = _optional_static_shape_list(target_shape)\\n"\n        "    if x.ndim != y.ndim:\\n"\n        "        return x, y\\n"\n        "    has_x_shape, x_shape = _tensor_static_shape_list(x)\\n"\n        "    has_y_shape, y_shape = _tensor_static_shape_list(y)\\n"\n        "    if not has_x_shape or not has_y_shape:\\n"\n        "        return x, y\\n"\n        "    if x_shape == y_shape:\\n"\n        "        return x, y\\n"\n        "    broadcast_shape = _broadcast_shape(x_shape, y_shape)\\n"\n        "    if broadcast_shape is not None:\\n"\n        "        x_is_scalar_seed = all(int(dim) == 1 for dim in list(x_shape))\\n"\n        "        y_is_scalar_seed = all(int(dim) == 1 for dim in list(y_shape))\\n"\n        "        if x_is_scalar_seed or y_is_scalar_seed:\\n"\n        "            return x, y\\n"\n        "        if not has_target or _shape_matches_target_relaxed(broadcast_shape, target):\\n"\n        "            return x, y\\n"\n        "        if torch.jit.is_scripting():\\n"\n        "            return x, y\\n"\n        "        return _align_binary_inputs_eager(x, y, target_shape)\\n"\n        "    if _can_broadcast_shapes(x_shape, y_shape):\\n"\n            "        if not has_target:\\n"\n            "            return x, y\\n"\n        "        broadcast_shape = _broadcast_shape(x_shape, y_shape)\\n"\n        "        if broadcast_shape is not None and _shape_matches_target_relaxed(broadcast_shape, target):\\n"\n        "                return x, y\\n"\n        "        if torch.jit.is_scripting():\\n"\n        "            return x, y\\n"\n        "        return _align_binary_inputs_eager(x, y, target_shape)\\n"\n        "    perm = _perm_cl_to_cf(x.ndim)\\n"\n        "    if perm is None:\\n"\n        "        return x, y\\n"\n        "    if _permute_shape(y_shape, perm) == x_shape:\\n"\n        "        return _align_binary_inputs_to_anchor(x, _torch_permute(y, perm), target_shape)\\n"\n        "    if _permute_shape(x_shape, perm) == y_shape:\\n"\n        "        return _align_binary_inputs_to_anchor(_torch_permute(x, perm), y, target_shape)\\n"\n        "    if has_target:\\n"\n        "        if _shape_matches_target_relaxed(_permute_shape(y_shape, perm), target):\\n"\n        "            return _align_binary_inputs_to_anchor(x, _torch_permute(y, perm), target_shape)\\n"\n        "        if _shape_matches_target_relaxed(_permute_shape(x_shape, perm), target):\\n"\n        "            return _align_binary_inputs_to_anchor(_torch_permute(x, perm), y, target_shape)\\n"\n        "    if torch.jit.is_scripting():\\n"\n        "        return x, y\\n"\n        "    return _align_binary_inputs_eager(x, y, target_shape)\\n\\n"\n        "@torch.jit.ignore\\n"\n        "def _align_binary_inputs_to_anchor_eager(anchor: torch.Tensor, other: torch.Tensor, target_shape: Optional[Sequence[int]]) -> Tuple[torch.Tensor, torch.Tensor]:\\n"\n        "    target = [int(v) for v in list(target_shape)] if target_shape is not None else None\\n"\n        "    anchor_shape = [int(v) for v in list(anchor.shape)]\\n"\n        "    other_shape = [int(v) for v in list(other.shape)]\\n"\n        "    perm = _perm_cl_to_cf(other.ndim)\\n"\n        "    if perm is not None and _permute_shape(other_shape, perm) == anchor_shape:\\n"\n        "        return anchor, _torch_permute(other, perm)\\n"\n        "    perm_inv = _perm_cf_to_cl(other.ndim)\\n"\n        "    if perm_inv is not None and _permute_shape(other_shape, perm_inv) == anchor_shape:\\n"\n        "        return anchor, _torch_permute(other, perm_inv)\\n"\n        "    if target is not None and other.ndim <= 5:\\n"\n        "        import itertools\\n"\n        "        for generic_perm in itertools.permutations(range(other.ndim)):\\n"\n        "            if list(generic_perm) == list(range(other.ndim)):\\n"\n        "                continue\\n"\n        "            generic_perm_list = [int(v) for v in list(generic_perm)]\\n"\n        "            permuted_shape = _permute_shape(other_shape, generic_perm_list)\\n"\n        "            if permuted_shape is None:\\n"\n        "                continue\\n"\n        "            try:\\n"\n        "                broadcast_shape = list(torch.broadcast_shapes(tuple(anchor_shape), tuple(permuted_shape)))\\n"\n        "                if _target_shape_has_concrete_dim(target) and _shape_matches_target_relaxed(broadcast_shape, target):\\n"\n        "                    return anchor, _torch_permute(other, generic_perm_list)\\n"\n        "            except Exception:\\n"\n        "                pass\\n"\n        "    if other.ndim <= 5:\\n"\n        "        import itertools\\n"\n        "        for generic_perm in itertools.permutations(range(other.ndim)):\\n"\n        "            if list(generic_perm) == list(range(other.ndim)):\\n"\n        "                continue\\n"\n        "            generic_perm_list = [int(v) for v in list(generic_perm)]\\n"\n        "            permuted_shape = _permute_shape(other_shape, generic_perm_list)\\n"\n        "            if permuted_shape == anchor_shape:\\n"\n        "                return anchor, _torch_permute(other, generic_perm_list)\\n"\n        "            if permuted_shape is not None:\\n"\n        "                try:\\n"\n        "                    broadcast_shape = list(torch.broadcast_shapes(tuple(anchor_shape), tuple(permuted_shape)))\\n"\n        "                    if broadcast_shape == anchor_shape or (target is not None and _target_shape_has_concrete_dim(target) and _shape_matches_target_relaxed(broadcast_shape, target)):\\n"\n        "                        return anchor, _torch_permute(other, generic_perm_list)\\n"\n        "                except Exception:\\n"\n        "                    pass\\n"\n        "    return anchor, other\\n\\n"\n        "def _align_binary_inputs_to_anchor(anchor: torch.Tensor, other: torch.Tensor, target_shape: Optional[Sequence[int]]) -> Tuple[torch.Tensor, torch.Tensor]:\\n"\n        "    has_target, target = _optional_static_shape_list(target_shape)\\n"\n        "    if anchor.ndim != other.ndim:\\n"\n        "        return anchor, other\\n"\n        "    has_anchor_shape, anchor_shape = _tensor_static_shape_list(anchor)\\n"\n        "    has_other_shape, other_shape = _tensor_static_shape_list(other)\\n"\n        "    if not has_anchor_shape or not has_other_shape:\\n"\n        "        return anchor, other\\n"\n        "    if anchor_shape == other_shape:\\n"\n        "        return anchor, other\\n"\n        "    broadcast_shape = _broadcast_shape(anchor_shape, other_shape)\\n"\n        "    if broadcast_shape is not None:\\n"\n        "        if broadcast_shape == anchor_shape:\\n"\n            "            return anchor, other\\n"\n        "        if not has_target:\\n"\n            "            return anchor, other\\n"\n        "        if _target_shape_has_concrete_dim(target) and _shape_matches_target_relaxed(broadcast_shape, target):\\n"\n            "            return anchor, other\\n"\n        "        if torch.jit.is_scripting():\\n"\n            "            return anchor, other\\n"\n        "        return _align_binary_inputs_to_anchor_eager(anchor, other, target_shape)\\n"\n        "    perm = _perm_cl_to_cf(other.ndim)\\n"\n        "    if perm is not None and _permute_shape(other_shape, perm) == anchor_shape:\\n"\n        "        return anchor, _torch_permute(other, perm)\\n"\n        "    perm_inv = _perm_cf_to_cl(other.ndim)\\n"\n        "    if perm_inv is not None and _permute_shape(other_shape, perm_inv) == anchor_shape:\\n"\n        "        return anchor, _torch_permute(other, perm_inv)\\n"\n        "    if has_target and _can_broadcast_shapes(anchor_shape, other_shape):\\n"\n        "        broadcast_shape = _broadcast_shape(anchor_shape, other_shape)\\n"\n        "        if broadcast_shape is not None and (broadcast_shape == anchor_shape or (_target_shape_has_concrete_dim(target) and _shape_matches_target_relaxed(broadcast_shape, target))):\\n"\n        "            return anchor, other\\n"\n        "        if torch.jit.is_scripting():\\n"\n        "            return anchor, other\\n"\n        "        return _align_binary_inputs_to_anchor_eager(anchor, other, target_shape)\\n"\n        "    if torch.jit.is_scripting():\\n"\n        "        return anchor, other\\n"\n        "    return _align_binary_inputs_to_anchor_eager(anchor, other, target_shape)\\n\\n"\n        "def _normalize_dim(dim: int, rank: int) -> int:\\n"\n        "    resolved = int(dim)\\n"\n        "    if resolved < 0:\\n"\n        "        resolved += int(rank)\\n"\n        "    return resolved\\n\\n"\n        "def _coerce_scalar_axis(value: Any, *, device: torch.device) -> int:\\n"\n        "    if isinstance(value, torch.Tensor):\\n"\n        "        flat = value.to(dtype=torch.int64, device=device).reshape(-1)\\n"\n        "        if int(flat.numel()) == 0:\\n"\n        "            return 0\\n"\n        "        return int(flat[0].item())\\n"\n        "    return int(value)\\n\\n"\n        "def _apply_cumsum(x: torch.Tensor, *, axis: int, exclusive: bool, reverse: bool) -> torch.Tensor:\\n"\n        "    dim = _normalize_dim(int(axis), x.ndim)\\n"\n        "    y = torch.flip(x, dims=[dim]) if reverse else x\\n"\n        "    y = torch.cumsum(y, dim=dim)\\n"\n        "    if exclusive:\\n"\n        "        axis_size = int(y.shape[dim])\\n"\n        "        if axis_size > 0:\\n"\n        "            zeros = torch.zeros_like(torch.narrow(y, dim, 0, 1))\\n"\n        "            prefix = torch.narrow(y, dim, 0, max(axis_size - 1, 0))\\n"\n        "            y = torch.cat([zeros, prefix], dim=dim)\\n"\n        "    if reverse:\\n"\n        "        y = torch.flip(y, dims=[dim])\\n"\n        "    return y\\n\\n"\n        "def _shape_list(value: Any) -> List[int]:\\n"\n        "    if isinstance(value, torch.Tensor):\\n"\n        "        flat = value.to(dtype=torch.int64).reshape(-1)\\n"\n        "        if torch.jit.is_scripting():\\n"\n        "            result = torch.jit.annotate(List[int], [])\\n"\n        "            for idx in range(int(flat.numel())):\\n"\n        "                result.append(int(flat[idx].item()))\\n"\n        "            return result\\n"\n        "        eager_result = []\\n"\n        "        for idx in range(int(flat.numel())):\\n"\n        "            eager_result.append(flat[idx])\\n"\n        "        return eager_result\\n"\n        "    if torch.jit.isinstance(value, List[int]):\\n"\n        "        result = torch.jit.annotate(List[int], [])\\n"\n        "        for item in value:\\n"\n        "            result.append(int(item))\\n"\n        "        return result\\n"\n        "    raise RuntimeError(\'Unsupported shape spec type for _shape_list\')\\n\\n"\n        "def _reshape_gather_output(gathered: torch.Tensor, params: torch.Tensor, indices_shape_spec: Sequence[int], *, axis: int) -> torch.Tensor:\\n"\n        "    resolved_axis = _normalize_dim(int(axis), params.ndim)\\n"\n        "    target_shape = torch.jit.annotate(List[int], [])\\n"\n        "    for dim_index in range(int(resolved_axis)):\\n"\n        "        target_shape.append(params.shape[dim_index])\\n"\n        "    for dim_value in indices_shape_spec:\\n"\n        "        target_shape.append(dim_value)\\n"\n        "    for dim_index in range(int(resolved_axis) + 1, int(params.ndim)):\\n"\n        "        target_shape.append(params.shape[dim_index])\\n"\n        "    return torch.reshape(gathered, target_shape)\\n\\n"\n        "@torch.jit.ignore\\n"\n        "def _apply_tile_eager(x: torch.Tensor, multiples: Any) -> torch.Tensor:\\n"\n        "    repeats = list(_shape_list(multiples))\\n"\n        "    target_rank = len(repeats)\\n"\n        "    if target_rank < int(x.ndim):\\n"\n        "        repeats = ([1] * (int(x.ndim) - target_rank)) + repeats\\n"\n        "        target_rank = len(repeats)\\n"\n        "    y = x\\n"\n        "    if target_rank > int(y.ndim):\\n"\n        "        reshape_shape = [1] * (target_rank - int(y.ndim))\\n"\n        "        for axis in range(int(y.ndim)):\\n"\n        "            reshape_shape.append(y.shape[axis])\\n"\n        "        y = torch.reshape(y, reshape_shape)\\n"\n        "    for dim_index in range(target_rank - 1, -1, -1):\\n"\n        "        repeat = repeats[dim_index]\\n"\n        "        y = torch.unsqueeze(y, dim_index)\\n"\n        "        expanded_shape = []\\n"\n        "        for axis in range(int(y.ndim)):\\n"\n        "            if axis == dim_index:\\n"\n        "                expanded_shape.append(repeat)\\n"\n        "                continue\\n"\n        "            expanded_shape.append(y.shape[axis])\\n"\n        "        y = y.expand(expanded_shape)\\n"\n        "        merged_shape = []\\n"\n        "        for axis in range(int(y.ndim)):\\n"\n        "            if axis == dim_index:\\n"\n        "                merged_shape.append(y.shape[axis] * y.shape[axis + 1])\\n"\n                "                continue\\n"\n        "            if axis == dim_index + 1:\\n"\n        "                continue\\n"\n        "            merged_shape.append(y.shape[axis])\\n"\n        "        y = torch.reshape(y, merged_shape)\\n"\n        "    return y\\n\\n"\n        "def _apply_tile(x: torch.Tensor, multiples: Any) -> torch.Tensor:\\n"\n        "    if torch.jit.is_scripting():\\n"\n        "        return torch.tile(x, _shape_list(multiples))\\n"\n        "    return _apply_tile_eager(x, multiples)\\n\\n"\n        "def _apply_random_standard_normal(shape_spec: Any, *, dtype: torch.dtype, device: torch.device, seed: Optional[int] = None) -> torch.Tensor:\\n"\n        "    shape = _shape_list(shape_spec)\\n"\n        "    if seed is None:\\n"\n        "        return torch.randn(shape, dtype=dtype, device=device)\\n"\n        "    if torch.jit.is_scripting():\\n"\n        "        torch.manual_seed(int(seed))\\n"\n        "        return torch.randn(shape, dtype=dtype, device=device)\\n"\n        "    generator = torch.Generator(device=device)\\n"\n        "    generator.manual_seed(int(seed))\\n"\n        "    return torch.randn(shape, dtype=dtype, device=device, generator=generator)\\n\\n"\n        "def _shape_tensor(input_tensor: torch.Tensor, *, dtype: torch.dtype, device: torch.device) -> torch.Tensor:\\n"\n        "    values = torch.jit.annotate(List[int], [])\\n"\n        "    for idx in range(int(input_tensor.ndim)):\\n"\n        "        values.append(input_tensor.shape[idx])\\n"\n        "    if len(values) == 0:\\n"\n        "        return torch.empty((0,), dtype=dtype, device=device)\\n"\n        "    return torch.tensor(values, dtype=dtype, device=device)\\n\\n"\n        "def _resolve_reshape_shape(shape_spec: Any, input_tensor: torch.Tensor, *, allow_zero: bool) -> List[int]:\\n"\n        "    raw_shape = _shape_list(shape_spec)\\n"\n        "    resolved: List[int] = []\\n"\n        "    raw_minus_one_count = 0\\n"\n        "    for raw_dim in raw_shape:\\n"\n        "        if int(raw_dim) == -1:\\n"\n        "            raw_minus_one_count += 1\\n"\n        "    infer_count = 0\\n"\n        "    input_rank = int(input_tensor.ndim)\\n"\n        "    for dim_index, raw_dim in enumerate(raw_shape):\\n"\n        "        dim_value = int(raw_dim)\\n"\n        "        if dim_index < input_rank and dim_value <= 0:\\n"\n        "            if dim_value == -1 and raw_minus_one_count > 1:\\n"\n        "                resolved.append(input_tensor.shape[dim_index])\\n"\n        "                continue\\n"\n        "            if dim_value == 0 and not allow_zero:\\n"\n        "                resolved.append(input_tensor.shape[dim_index])\\n"\n        "                continue\\n"\n        "        if dim_value == 0 and not allow_zero and dim_index < input_rank:\\n"\n        "            resolved.append(input_tensor.shape[dim_index])\\n"\n        "            continue\\n"\n        "        if dim_value == -1:\\n"\n        "            infer_count += 1\\n"\n        "            if infer_count > 1:\\n"\n        "                raise RuntimeError(f\'Multiple -1 values are not allowed in reshape spec: {raw_shape}\')\\n"\n        "        resolved.append(dim_value)\\n"\n        "    if infer_count == 1:\\n"\n        "        infer_index = -1\\n"\n        "        known_product = 1\\n"\n        "        for dim_index, dim_value in enumerate(resolved):\\n"\n        "            if int(dim_value) == -1:\\n"\n        "                infer_index = int(dim_index)\\n"\n        "                continue\\n"\n        "            known_product *= dim_value\\n"\n        "        if infer_index >= 0:\\n"\n        "            resolved[infer_index] = input_tensor.numel() // known_product\\n"\n        "    return resolved\\n\\n"\n        "def _resolve_reshape_shape_tensor(shape_spec: torch.Tensor, input_tensor: torch.Tensor, *, allow_zero: bool) -> torch.Tensor:\\n"\n        "    flat = shape_spec.to(dtype=torch.int64, device=input_tensor.device).reshape(-1)\\n"\n        "    if allow_zero:\\n"\n        "        return flat\\n"\n        "    resolved = torch.jit.annotate(List[torch.Tensor], [])\\n"\n        "    input_rank = int(input_tensor.ndim)\\n"\n        "    for dim_index in range(int(flat.numel())):\\n"\n        "        dim_value = flat[dim_index]\\n"\n        "        if dim_index < input_rank:\\n"\n        "            input_dim = torch.scalar_tensor(input_tensor.shape[dim_index], dtype=flat.dtype, device=flat.device)\\n"\n        "            resolved.append(torch.where(dim_value == 0, input_dim, dim_value))\\n"\n        "            continue\\n"\n        "        resolved.append(dim_value)\\n"\n        "    if len(resolved) == 0:\\n"\n        "        return flat\\n"\n        "    return torch.stack(resolved)\\n\\n"\n        "def _to_torch_pad_arg(paddings: torch.Tensor) -> List[int]:\\n"\n        "    pads_tensor = paddings.to(dtype=torch.int64).reshape(-1, 2)\\n"\n        "    torch_pad: List[int] = []\\n"\n        "    for row_index in range(int(pads_tensor.shape[0]) - 1, -1, -1):\\n"\n        "        torch_pad.extend([\\n"\n        "            int(pads_tensor[row_index, 0].item()),\\n"\n        "            int(pads_tensor[row_index, 1].item()),\\n"\n        "        ])\\n"\n        "    while len(torch_pad) >= 2 and int(torch_pad[-2]) == 0 and int(torch_pad[-1]) == 0:\\n"\n        "        torch_pad = torch_pad[:-2]\\n"\n        "    return torch_pad\\n\\n"\n        "def _apply_pad_nd(x: torch.Tensor, paddings: torch.Tensor, *, mode: str, value: float = 0.0) -> torch.Tensor:\\n"\n        "    pads_tensor = paddings.to(dtype=torch.int64).reshape(-1, 2)\\n"\n        "    if (not torch.jit.is_scripting()) and str(mode) == \'constant\':\\n"\n        "        rank = int(x.ndim)\\n"\n        "        pad_rows = int(pads_tensor.shape[0])\\n"\n        "        if pad_rows < rank:\\n"\n        "            prefix_rows = torch.zeros([rank - pad_rows, 2], dtype=pads_tensor.dtype, device=pads_tensor.device)\\n"\n        "            pads_tensor = torch.cat([prefix_rows, pads_tensor], dim=0)\\n"\n        "        elif pad_rows > rank:\\n"\n        "            pads_tensor = pads_tensor[(pad_rows - rank):, :]\\n"\n        "        y = x\\n"\n        "        for axis in range(rank):\\n"\n        "            before = torch.reshape(pads_tensor[axis, 0], [1]).to(dtype=torch.int64)\\n"\n        "            after = torch.reshape(pads_tensor[axis, 1], [1]).to(dtype=torch.int64)\\n"\n        "            zero_shape = list(y.shape)\\n"\n        "            zero_shape[axis] = 1\\n"\n        "            zero_slice = torch.full(zero_shape, float(value), dtype=y.dtype, device=y.device)\\n"\n        "            pad_before = torch.repeat_interleave(zero_slice, before, dim=axis)\\n"\n        "            pad_after = torch.repeat_interleave(zero_slice, after, dim=axis)\\n"\n        "            y = torch.cat([pad_before, y, pad_after], dim=axis)\\n"\n        "        return y\\n"\n        "    pad_pairs = torch.jit.annotate(List[List[int]], [])\\n"\n        "    for row_index in range(int(pads_tensor.shape[0])):\\n"\n        "        pad_pairs.append([\\n"\n        "            int(pads_tensor[row_index, 0].item()),\\n"\n        "            int(pads_tensor[row_index, 1].item()),\\n"\n        "        ])\\n"\n        "    rank = int(x.ndim)\\n"\n        "    if len(pad_pairs) < rank:\\n"\n        "        pad_pairs = ([[0, 0]] * (rank - len(pad_pairs))) + pad_pairs\\n"\n        "    elif len(pad_pairs) > rank:\\n"\n        "        pad_pairs = pad_pairs[-rank:]\\n"\n        "    non_zero_axes = [idx for idx, (before, after) in enumerate(pad_pairs) if int(before) != 0 or int(after) != 0]\\n"\n        "    if len(non_zero_axes) == 0:\\n"\n        "        return x\\n"\n        "    if mode != \'constant\' and len(non_zero_axes) > 3:\\n"\n        "        raise RuntimeError(f\'Non-constant pad supports at most 3 padded dims. mode={mode} padded_dims={len(non_zero_axes)}\')\\n"\n        "    keep_axes = [idx for idx in range(rank) if idx not in non_zero_axes]\\n"\n        "    perm = keep_axes + non_zero_axes\\n"\n        "    permuted = _torch_permute(x, perm) if perm != list(range(rank)) else x\\n"\n        "    torch_pad: List[int] = []\\n"\n        "    for axis in reversed(non_zero_axes):\\n"\n        "        before, after = pad_pairs[axis]\\n"\n        "        torch_pad.extend([int(before), int(after)])\\n"\n        "    if mode == \'constant\':\\n"\n        "        padded = F.pad(permuted, torch_pad, mode=mode, value=float(value))\\n"\n        "    else:\\n"\n        "        padded = F.pad(permuted, torch_pad, mode=mode)\\n"\n        "    if perm == list(range(rank)):\\n"\n        "        return padded\\n"\n        "    inverse_perm = [0] * rank\\n"\n        "    for permuted_axis, original_axis in enumerate(perm):\\n"\n        "        inverse_perm[int(original_axis)] = int(permuted_axis)\\n"\n        "    return _torch_permute(padded, inverse_perm)\\n\\n"\n        "def _infer_spatial_shape_for_transposed_conv2d(*, raw_output: torch.Tensor, target_shape: Optional[Sequence[int]], fallback_shape: Sequence[int]) -> Tuple[int, int]:\\n"\n        "    output_channels = int(raw_output.shape[1])\\n"\n        "    source = [int(v) for v in list(target_shape)] if target_shape is not None else [int(v) for v in list(fallback_shape)]\\n"\n        "    if len(source) == 4:\\n"\n        "        if int(source[1]) == output_channels:\\n"\n        "            return int(source[2]), int(source[3])\\n"\n        "        if int(source[-1]) == output_channels:\\n"\n        "            return int(source[1]), int(source[2])\\n"\n        "    return int(source[-2]), int(source[-1])\\n\\n"\n        "def _infer_spatial_shape_for_transposed_conv3d(*, raw_output: torch.Tensor, target_shape: Optional[Sequence[int]], fallback_shape: Sequence[int]) -> Tuple[int, int, int]:\\n"\n        "    output_channels = int(raw_output.shape[1])\\n"\n        "    source = [int(v) for v in list(target_shape)] if target_shape is not None else [int(v) for v in list(fallback_shape)]\\n"\n        "    if len(source) == 5:\\n"\n        "        if int(source[1]) == output_channels:\\n"\n        "            return int(source[2]), int(source[3]), int(source[4])\\n"\n        "        if int(source[-1]) == output_channels:\\n"\n        "            return int(source[1]), int(source[2]), int(source[3])\\n"\n        "    return int(source[-3]), int(source[-2]), int(source[-1])\\n\\n"\n        "def _apply_fused_activation(x: torch.Tensor, fused: str) -> torch.Tensor:\\n"\n        "    key = str(fused).upper()\\n"\n        "    if key == \'\' or key == \'NONE\':\\n"\n        "        return x\\n"\n        "    if key == \'RELU\':\\n"\n        "        return torch.relu(x)\\n"\n        "    if key == \'RELU6\':\\n"\n        "        return torch.clamp(x, min=0.0, max=6.0)\\n"\n        "    if key == \'RELU_N1_TO_1\':\\n"\n        "        return torch.clamp(x, min=-1.0, max=1.0)\\n"\n        "    if key == \'RELU_0_TO_1\':\\n"\n        "        return torch.clamp(x, min=0.0, max=1.0)\\n"\n        "    if key == \'SILU\':\\n"\n        "        return torch.mul(x, torch.sigmoid(x))\\n"\n        "    if key == \'TANH\':\\n"\n        "        return torch.tanh(x)\\n"\n        "    return x\\n\\n"\n        "def _lookup_state_tensor(raw_state_dict: Dict[str, Any], tensor_name: str, storage_names: Dict[str, str]) -> torch.Tensor:\\n"\n        "    original_key = str(tensor_name)\\n"\n        "    storage_key = storage_names.get(original_key, _default_tensor_storage_name(original_key))\\n"\n        "    if original_key in raw_state_dict:\\n"\n        "        return torch.as_tensor(raw_state_dict[original_key])\\n"\n        "    if storage_key in raw_state_dict:\\n"\n        "        return torch.as_tensor(raw_state_dict[storage_key])\\n"\n        "    raise KeyError(original_key)\\n\\n"\n        "def _copy_tensor_data(target: torch.Tensor, source: torch.Tensor) -> None:\\n"\n        "    target.data.copy_(source.to(device=target.device, dtype=target.dtype))\\n\\n"\n        "def _validate_state_dict_keys(raw_state_dict: Dict[str, Any], storage_names: Dict[str, str], expected_tensor_names: Sequence[str]) -> None:\\n"\n        "    recognized_keys: Set[str] = set()\\n"\n        "    missing: List[str] = []\\n"\n        "    for tensor_name in expected_tensor_names:\\n"\n        "        storage_key = storage_names.get(str(tensor_name), _default_tensor_storage_name(str(tensor_name)))\\n"\n        "        if str(tensor_name) in raw_state_dict:\\n"\n        "            recognized_keys.add(str(tensor_name))\\n"\n        "            continue\\n"\n        "        if storage_key in raw_state_dict:\\n"\n        "            recognized_keys.add(storage_key)\\n"\n        "            continue\\n"\n        "        missing.append(str(tensor_name))\\n"\n        "    unexpected = sorted(str(key) for key in raw_state_dict.keys() if str(key) not in recognized_keys)\\n"\n        "    if len(missing) > 0 or len(unexpected) > 0:\\n"\n        "        raise RuntimeError(f\'state_dict mismatch. missing={missing} unexpected={unexpected}\')\\n\\n"\n        "def _shape_list_contains(candidates: List[List[int]], target: List[int]) -> bool:\\n"\n        "    for candidate in candidates:\\n"\n        "        if len(candidate) != len(target):\\n"\n        "            continue\\n"\n        "        matches = True\\n"\n        "        for index in range(len(target)):\\n"\n        "            if int(candidate[index]) != int(target[index]):\\n"\n        "                matches = False\\n"\n        "                break\\n"\n        "        if matches:\\n"\n        "            return True\\n"\n        "    return False\\n\\n"\n        "def _apply_concat(values: Sequence[torch.Tensor], axis: int, target_shape: Optional[Sequence[int]], fused: str) -> torch.Tensor:\\n"\n        "    if any(int(value.ndim) == 0 for value in values):\\n"\n        "        values = [value.reshape(1) if int(value.ndim) == 0 else value for value in values]\\n"\n        "    rank = int(values[0].ndim)\\n"\n        "    resolved_axis = _normalize_dim(int(axis), rank)\\n"\n        "    has_target, target_list = _optional_static_shape_list(target_shape)\\n"\n        "    if has_target and len(target_list) == rank:\\n"\n        "        candidate_targets: List[List[int]] = [target_list]\\n"\n        "        perm_cf_to_cl = _perm_cf_to_cl(rank)\\n"\n        "        if perm_cf_to_cl is not None:\\n"\n        "            alt_target = _permute_shape(target_list, perm_cf_to_cl)\\n"\n        "            if alt_target is not None and not _shape_list_contains(candidate_targets, alt_target):\\n"\n        "                candidate_targets.append(alt_target)\\n"\n        "        perm_cl_to_cf = _perm_cl_to_cf(rank)\\n"\n        "        if perm_cl_to_cf is not None:\\n"\n        "            alt_target = _permute_shape(target_list, perm_cl_to_cf)\\n"\n        "            if alt_target is not None and not _shape_list_contains(candidate_targets, alt_target):\\n"\n        "                candidate_targets.append(alt_target)\\n"\n        "        best_aligned_values: Optional[List[torch.Tensor]] = None\\n"\n        "        best_score: Optional[int] = None\\n"\n        "        for candidate_target in candidate_targets:\\n"\n        "            aligned_values = []\\n"\n        "            candidate_score = 0\\n"\n        "            feasible = True\\n"\n        "            for value in values:\\n"\n        "                has_actual, actual = _tensor_static_shape_list(value)\\n"\n        "                if not has_actual:\\n"\n        "                    aligned_values.append(value)\\n"\n        "                    continue\\n"\n        "                chosen = value\\n"\n        "                if actual != candidate_target:\\n"\n        "                    if _matches_target_except_axis(actual, candidate_target, resolved_axis):\\n"\n        "                        aligned_values.append(chosen)\\n"\n        "                        continue\\n"\n        "                    matched = False\\n"\n        "                    perm_cl_to_cf_value = _perm_cl_to_cf(value.ndim)\\n"\n        "                    if perm_cl_to_cf_value is not None:\\n"\n        "                        permuted_shape = _permute_shape(actual, perm_cl_to_cf_value)\\n"\n        "                        if _matches_target_except_axis(permuted_shape, candidate_target, resolved_axis):\\n"\n        "                            chosen = _torch_permute(value, perm_cl_to_cf_value)\\n"\n        "                            candidate_score += 1\\n"\n        "                            matched = True\\n"\n        "                    if not matched:\\n"\n        "                        perm_cf_to_cl_value = _perm_cf_to_cl(value.ndim)\\n"\n        "                        if perm_cf_to_cl_value is not None:\\n"\n        "                            permuted_shape = _permute_shape(actual, perm_cf_to_cl_value)\\n"\n        "                            if _matches_target_except_axis(permuted_shape, candidate_target, resolved_axis):\\n"\n        "                                chosen = _torch_permute(value, perm_cf_to_cl_value)\\n"\n        "                                candidate_score += 1\\n"\n        "                                matched = True\\n"\n        "                    if not matched:\\n"\n        "                        feasible = False\\n"\n        "                        break\\n"\n        "                aligned_values.append(chosen)\\n"\n        "            if feasible and (best_score is None or int(candidate_score) < int(best_score)):\\n"\n        "                best_aligned_values = aligned_values\\n"\n        "                best_score = int(candidate_score)\\n"\n        "        if best_aligned_values is not None:\\n"\n        "            values = best_aligned_values\\n"\n        "    y = torch.cat(list(values), dim=resolved_axis)\\n"\n        "    y = _align_tensor_to_target_shape(y, target_shape)\\n"\n        "    return _apply_fused_activation(y, fused)\\n\\n"\n        "def _apply_module_conv2d(module: torch.nn.Conv2d, x: torch.Tensor, target_shape: Optional[Sequence[int]], target_logical_layout: Optional[str], fused: str) -> torch.Tensor:\\n"\n        "    expected_in_channels = int(module.in_channels)\\n"\n        "    if x.ndim == 4 and int(x.shape[1]) != expected_in_channels and int(x.shape[-1]) == expected_in_channels:\\n"\n        "        x = x.permute(0, 3, 1, 2).contiguous()\\n"\n        "    y = module.forward(x)\\n"\n        "    if str(target_logical_layout).upper() == \'NHWC\':\\n"\n        "        y = y.permute(0, 2, 3, 1).contiguous()\\n"\n        "    y = _align_tensor_to_target_shape(y, target_shape)\\n"\n        "    return _apply_fused_activation(y, fused)\\n\\n"\n        "def _apply_module_transpose_conv2d(x: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor], stride: List[int], padding: List[int], dilation: List[int], output_padding: List[int], groups: int, target_shape: Optional[Sequence[int]], fallback_shape: Sequence[int], target_logical_layout: Optional[str], fused: str) -> torch.Tensor:\\n"\n        "    if x.ndim == 4 and int(x.shape[1]) != int(weight.shape[0]) and int(x.shape[-1]) == int(weight.shape[0]):\\n"\n        "        x = x.permute(0, 3, 1, 2).contiguous()\\n"\n        "    raw = F.conv_transpose2d(x, weight, bias=bias, stride=stride, padding=padding, output_padding=output_padding, groups=int(groups), dilation=dilation)\\n"\n        "    target_h, target_w = _infer_spatial_shape_for_transposed_conv2d(raw_output=raw, target_shape=target_shape, fallback_shape=fallback_shape)\\n"\n        "    y = raw[..., :target_h, :target_w]\\n"\n        "    if str(target_logical_layout).upper() == \'NHWC\':\\n"\n        "        y = y.permute(0, 2, 3, 1).contiguous()\\n"\n        "    y = _align_tensor_to_target_shape(y, target_shape)\\n"\n        "    return _apply_fused_activation(y, fused)\\n\\n"\n        "def _apply_module_conv3d(module: torch.nn.Conv3d, x: torch.Tensor, target_shape: Optional[Sequence[int]], target_logical_layout: Optional[str], fused: str) -> torch.Tensor:\\n"\n        "    weight = module.weight\\n"\n        "    if x.ndim == 5 and int(x.shape[1]) != int(weight.shape[1]) and int(x.shape[-1]) == int(weight.shape[1]):\\n"\n        "        x = x.permute(0, 4, 1, 2, 3).contiguous()\\n"\n        "    y = module.forward(x)\\n"\n        "    if str(target_logical_layout).upper() == \'NDHWC\':\\n"\n        "        y = y.permute(0, 2, 3, 4, 1).contiguous()\\n"\n        "    y = _align_tensor_to_target_shape(y, target_shape)\\n"\n        "    return _apply_fused_activation(y, fused)\\n\\n"\n        "def _apply_module_transpose_conv3d(x: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor], stride: List[int], padding: List[int], dilation: List[int], output_padding: List[int], groups: int, target_shape: Optional[Sequence[int]], fallback_shape: Sequence[int], target_logical_layout: Optional[str], fused: str) -> torch.Tensor:\\n"\n        "    if x.ndim == 5 and int(x.shape[1]) != int(weight.shape[0]) and int(x.shape[-1]) == int(weight.shape[0]):\\n"\n        "        x = x.permute(0, 4, 1, 2, 3).contiguous()\\n"\n        "    raw = F.conv_transpose3d(x, weight, bias=bias, stride=stride, padding=padding, output_padding=output_padding, groups=int(groups), dilation=dilation)\\n"\n        "    target_d, target_h, target_w = _infer_spatial_shape_for_transposed_conv3d(raw_output=raw, target_shape=target_shape, fallback_shape=fallback_shape)\\n"\n        "    y = raw[..., :target_d, :target_h, :target_w]\\n"\n        "    if str(target_logical_layout).upper() == \'NDHWC\':\\n"\n        "        y = y.permute(0, 2, 3, 4, 1).contiguous()\\n"\n        "    y = _align_tensor_to_target_shape(y, target_shape)\\n"\n        "    return _apply_fused_activation(y, fused)\\n\\n"\n        "def _apply_softmax(x: torch.Tensor, axis: Optional[int], beta: float, target_shape: Optional[Sequence[int]]) -> torch.Tensor:\\n"\n        "    resolved_axis = _normalize_dim(int(axis), x.ndim) if axis is not None else -1\\n"\n        "    if beta != 1.0:\\n"\n        "        x = x * beta\\n"\n        "    y = torch.softmax(x, dim=resolved_axis)\\n"\n        "    return _align_tensor_to_target_shape(y, target_shape)\\n\\n"\n        "@torch.jit.ignore\\n"\n        "def _apply_gather_batched_eager(params: torch.Tensor, indices_i64: torch.Tensor, *, resolved_axis: int, resolved_batch_dims: int, target_shape: Optional[Sequence[int]]) -> torch.Tensor:\\n"\n        "    params_flat = torch.flatten(params, 0, resolved_batch_dims - 1)\\n"\n        "    indices_flat = torch.flatten(indices_i64, 0, resolved_batch_dims - 1)\\n"\n        "    adjusted_axis = int(resolved_axis - resolved_batch_dims + 1)\\n"\n        "    prefix_shape = list(params_flat.shape[1:adjusted_axis])\\n"\n        "    suffix_shape = list(params_flat.shape[adjusted_axis + 1:])\\n"\n        "    index_shape = list(indices_flat.shape[1:])\\n"\n        "    input_view = params_flat\\n"\n        "    for _ in range(max(len(index_shape) - 1, 0)):\\n"\n        "        input_view = input_view.unsqueeze(adjusted_axis + 1)\\n"\n        "    input_expand_shape = list(params_flat.shape[:adjusted_axis + 1]) + list(index_shape[1:]) + suffix_shape\\n"\n        "    input_expanded = input_view.expand(input_expand_shape)\\n"\n        "    index_view_shape = [indices_flat.shape[0]] + ([1] * len(prefix_shape)) + index_shape + ([1] * len(suffix_shape))\\n"\n        "    index_view = indices_flat.reshape(index_view_shape)\\n"\n        "    index_expand_shape = [indices_flat.shape[0]] + prefix_shape + index_shape + suffix_shape\\n"\n        "    index_expanded = index_view.expand(index_expand_shape)\\n"\n        "    y = torch.gather(input_expanded, adjusted_axis, index_expanded)\\n"\n        "    return _align_tensor_to_target_shape(y, target_shape)\\n\\n"\n        "def _apply_gather(params: torch.Tensor, indices: torch.Tensor, axis: int, batch_dims: int, target_shape: Optional[Sequence[int]], indices_name: str) -> torch.Tensor:\\n"\n        "    indices_i64 = indices.to(dtype=torch.int64)\\n"\n        "    resolved_axis = _normalize_dim(int(axis), params.ndim)\\n"\n        "    if int(batch_dims) == 0 and int(resolved_axis) == 1 and str(indices_name).endswith(\'_crd_to_dcr_indices\'):\\n"\n        "        return _align_tensor_to_target_shape(params, target_shape)\\n"\n        "    resolved_batch_dims = int(batch_dims)\\n"\n        "    if resolved_batch_dims < 0:\\n"\n        "        resolved_batch_dims += indices_i64.ndim\\n"\n        "    if resolved_batch_dims > 0:\\n"\n        "        if not torch.jit.is_scripting():\\n"\n        "            return _apply_gather_batched_eager(params, indices_i64, resolved_axis=resolved_axis, resolved_batch_dims=resolved_batch_dims, target_shape=target_shape)\\n"\n        "        leading_shape = torch.jit.annotate(List[int], [])\\n"\n        "        for dim_index in range(int(resolved_batch_dims)):\\n"\n        "            leading_shape.append(indices_i64.shape[dim_index])\\n"\n        "        params_flat = torch.flatten(params, 0, resolved_batch_dims - 1)\\n"\n        "        indices_flat = torch.flatten(indices_i64, 0, resolved_batch_dims - 1)\\n"\n        "        gathered_batches: List[torch.Tensor] = []\\n"\n        "        adjusted_axis = int(resolved_axis - resolved_batch_dims + 1)\\n"\n        "        for batch_index in range(params_flat.shape[0]):\\n"\n        "            batch_params = params_flat[batch_index]\\n"\n        "            batch_indices = indices_flat[batch_index]\\n"\n        "            flat_indices = batch_indices.reshape(-1)\\n"\n        "            batch_gathered = torch.index_select(batch_params, adjusted_axis - 1, flat_indices)\\n"\n        "            batch_gathered = _reshape_gather_output(batch_gathered, batch_params, list(batch_indices.shape), axis=adjusted_axis - 1)\\n"\n        "            gathered_batches.append(batch_gathered)\\n"\n        "        batch_result_shape = torch.jit.annotate(List[int], [])\\n"\n        "        for dim_value in leading_shape:\\n"\n        "            batch_result_shape.append(dim_value)\\n"\n        "        for dim_value in gathered_batches[0].shape:\\n"\n        "            batch_result_shape.append(dim_value)\\n"\n        "        y = torch.reshape(torch.stack(gathered_batches, dim=0), batch_result_shape)\\n"\n        "        return _align_tensor_to_target_shape(y, target_shape)\\n"\n        "    if indices_i64.ndim == 0:\\n"\n        "        y = torch.index_select(params, resolved_axis, indices_i64.reshape(1)).squeeze(resolved_axis)\\n"\n        "        return _align_tensor_to_target_shape(y, target_shape)\\n"\n        "    flat_indices = indices_i64.reshape(-1)\\n"\n        "    gathered = torch.index_select(params, resolved_axis, flat_indices)\\n"\n        "    y = _reshape_gather_output(gathered, params, list(indices_i64.shape), axis=resolved_axis)\\n"\n        "    return _align_tensor_to_target_shape(y, target_shape)\\n\\n"\n        "def _apply_gather_elements(params: torch.Tensor, indices: torch.Tensor, axis_hint: int, axis_hint_cl: Optional[int], target_shape: Optional[Sequence[int]]) -> torch.Tensor:\\n"\n        "    indices_i64 = indices.to(dtype=torch.int64)\\n"\n        "    def _gather_for_axis(axis: int) -> Tuple[torch.Tensor, torch.Tensor]:\\n"\n        "        resolved_axis = int(axis)\\n"\n        "        dim_size = int(params.shape[resolved_axis])\\n"\n        "        if dim_size <= 0:\\n"\n        "            safe_indices = torch.zeros_like(indices_i64)\\n"\n        "            gathered = torch.gather(params, dim=resolved_axis, index=safe_indices)\\n"\n        "            return gathered, torch.zeros((), dtype=torch.bool, device=indices_i64.device)\\n"\n        "        dim_size_tensor = torch.as_tensor(dim_size, dtype=torch.int64, device=indices_i64.device)\\n"\n        "        normalized_indices = torch.where(indices_i64 < 0, indices_i64 + dim_size_tensor, indices_i64)\\n"\n        "        axis_valid = torch.all((indices_i64 >= -dim_size_tensor) & (indices_i64 < dim_size_tensor))\\n"\n        "        safe_indices = torch.clamp(normalized_indices, min=0, max=dim_size - 1)\\n"\n        "        return torch.gather(params, dim=resolved_axis, index=safe_indices), axis_valid\\n"\n        "    gather_axis = int(axis_hint)\\n"\n        "    y, y_valid = _gather_for_axis(gather_axis)\\n"\n        "    candidate_axes: List[int] = []\\n"\n        "    if axis_hint_cl is not None and int(axis_hint_cl) != gather_axis:\\n"\n        "        candidate_axes.append(int(axis_hint_cl))\\n"\n        "    candidate_axes.extend(\\n"\n        "        axis\\n"\n        "        for axis in range(params.ndim - 1, -1, -1)\\n"\n        "        if axis != gather_axis and (axis_hint_cl is None or axis != int(axis_hint_cl))\\n"\n        "    )\\n"\n        "    for candidate_axis in candidate_axes:\\n"\n        "        candidate_y, candidate_valid = _gather_for_axis(int(candidate_axis))\\n"\n        "        use_candidate = torch.logical_and(torch.logical_not(y_valid), candidate_valid)\\n"\n        "        y = torch.where(use_candidate, candidate_y, y)\\n"\n        "        y_valid = torch.logical_or(y_valid, candidate_valid)\\n"\n        "    return _align_tensor_to_target_shape(y, target_shape)\\n\\n"\n        "def _apply_gather_nd(params: torch.Tensor, indices: torch.Tensor, target_shape: Optional[Sequence[int]]) -> torch.Tensor:\\n"\n        "    indices_i64 = indices.to(dtype=torch.int64)\\n"\n        "    index_depth = int(indices_i64.shape[-1])\\n"\n        "    flat_indices = indices_i64.reshape(-1, index_depth)\\n"\n        "    prefix_shape = list(indices_i64.shape[:-1])\\n"\n        "    slice_shape = list(params.shape[index_depth:])\\n"\n        "    if index_depth == 0:\\n"\n        "        gathered_rows = params\\n"\n        "    elif not torch.jit.is_scripting():\\n"\n        "        gather_indices = []\\n"\n        "        for axis in range(index_depth):\\n"\n        "            gather_indices.append(flat_indices[:, axis])\\n"\n        "        gathered_rows = params[tuple(gather_indices)]\\n"\n        "    else:\\n"\n        "        row_shape = [int(v) for v in list(params.shape[:index_depth])]\\n"\n        "        strides = [1] * index_depth\\n"\n        "        running = 1\\n"\n        "        for axis in range(index_depth - 1, -1, -1):\\n"\n        "            strides[axis] = int(running)\\n"\n        "            running *= int(row_shape[axis])\\n"\n        "        flat_row_indices = torch.zeros_like(flat_indices[:, 0])\\n"\n        "        for axis in range(index_depth):\\n"\n        "            axis_indices = flat_indices[:, axis]\\n"\n        "            dim_size = int(row_shape[axis])\\n"\n        "            if dim_size > 0:\\n"\n        "                dim_size_tensor = torch.as_tensor(dim_size, dtype=torch.int64, device=indices_i64.device)\\n"\n        "                axis_indices = torch.where(axis_indices < 0, axis_indices + dim_size_tensor, axis_indices)\\n"\n        "            flat_row_indices = flat_row_indices + (axis_indices * int(strides[axis]))\\n"\n        "        flat_params_shape = torch.jit.annotate(List[int], [int(running)])\\n"\n        "        for dim_value in slice_shape:\\n"\n        "            flat_params_shape.append(int(dim_value))\\n"\n        "        flat_params = torch.reshape(params, flat_params_shape)\\n"\n        "        gathered_rows = torch.index_select(flat_params, 0, flat_row_indices)\\n"\n        "    output_shape = prefix_shape + slice_shape\\n"\n        "    if len(output_shape) == 0:\\n"\n        "        y = torch.reshape(gathered_rows, [])\\n"\n        "    else:\\n"\n        "        y = torch.reshape(gathered_rows, output_shape)\\n"\n        "    return _align_tensor_to_target_shape(y, target_shape)\\n\\n"\n        "def _apply_scatter_nd(indices: torch.Tensor, updates: torch.Tensor, shape: Any, target_shape: Optional[Sequence[int]]) -> torch.Tensor:\\n"\n        "    output_shape = _shape_list(shape)\\n"\n        "    indices_i64 = indices.to(dtype=torch.int64)\\n"\n        "    index_depth = int(indices_i64.shape[-1])\\n"\n        "    prefix_shape = list(indices_i64.shape[:-1])\\n"\n        "    slice_shape = list(output_shape[index_depth:])\\n"\n        "    expected_updates_shape = prefix_shape + slice_shape\\n"\n        "    aligned_updates = _align_scatter_nd_updates(updates, expected_updates_shape)\\n"\n        "    flat_indices = indices_i64.reshape(-1, index_depth)\\n"\n        "    row_shape = [int(v) for v in output_shape[:index_depth]]\\n"\n        "    strides: List[int] = [1] * len(row_shape)\\n"\n        "    row_count = 1\\n"\n        "    running = 1\\n"\n        "    for axis in range(len(row_shape) - 1, -1, -1):\\n"\n        "        strides[axis] = int(running)\\n"\n        "        running *= int(row_shape[axis])\\n"\n        "    row_count = int(running)\\n"\n        "    slice_size = 1\\n"\n        "    for dim in slice_shape:\\n"\n        "        slice_size *= int(dim)\\n"\n        "    stride_tensor = torch.as_tensor(strides, dtype=torch.int64, device=indices_i64.device)\\n"\n        "    flat_rows = torch.sum(flat_indices * stride_tensor, dim=-1)\\n"\n        "    if int(slice_size) == 1:\\n"\n        "        y_flat = torch.zeros((row_count,), dtype=aligned_updates.dtype, device=aligned_updates.device)\\n"\n        "        y_flat.index_copy_(0, flat_rows, aligned_updates.reshape(-1))\\n"\n        "        y = y_flat.reshape(output_shape)\\n"\n        "    else:\\n"\n        "        y_rows = torch.zeros((row_count, slice_size), dtype=aligned_updates.dtype, device=aligned_updates.device)\\n"\n        "        y_rows.index_copy_(0, flat_rows, aligned_updates.reshape(-1, slice_size))\\n"\n        "        y = y_rows.reshape(output_shape)\\n"\n        "    return _align_tensor_to_target_shape(y, target_shape)\\n\\n"\n        "@torch.jit.ignore\\n"\n        "def _apply_if_axis0_tensor_mux_export(cond: torch.Tensor, then_value: torch.Tensor, else_value: torch.Tensor) -> torch.Tensor:\\n"\n        "    predicate = torch.reshape(cond.to(dtype=torch.bool), [])\\n"\n        "    def _then(branch_then: torch.Tensor, branch_else: torch.Tensor) -> torch.Tensor:\\n"\n        "        return branch_then.clone()\\n"\n        "    def _else(branch_then: torch.Tensor, branch_else: torch.Tensor) -> torch.Tensor:\\n"\n        "        return branch_else.clone()\\n"\n        "    return torch.cond(predicate, _then, _else, (then_value, else_value))\\n\\n"\n        "def _apply_if_axis0_tensor_mux(cond: torch.Tensor, then_value: torch.Tensor, else_value: torch.Tensor, merged: torch.Tensor, begin: torch.Tensor, size: torch.Tensor, target_shape: Optional[Sequence[int]], use_export_mode: bool) -> torch.Tensor:\\n"\n        "    if bool(use_export_mode):\\n"\n        "        y = _apply_if_axis0_tensor_mux_export(cond, then_value, else_value)\\n"\n        "        return _align_tensor_to_target_shape(y, target_shape)\\n"\n        "    return _apply_slice(merged, begin, size, target_shape)\\n\\n"\n        "def _apply_slice(x: torch.Tensor, begin: torch.Tensor, size: torch.Tensor, target_shape: Optional[Sequence[int]]) -> torch.Tensor:\\n"\n        "    begin_values = _shape_list(begin)\\n"\n        "    size_values = _shape_list(size)\\n"\n        "    y = x\\n"\n        "    for axis, start in enumerate(begin_values):\\n"\n        "        dim_size = y.shape[axis]\\n"\n        "        start_index = int(start)\\n"\n        "        if start_index < 0:\\n"\n        "            start_index += dim_size\\n"\n        "        length = int(size_values[axis])\\n"\n        "        if length < 0:\\n"\n        "            narrowed_length = dim_size - start_index\\n"\n        "        else:\\n"\n        "            narrowed_length = length\\n"\n        "        y = torch.narrow(y, axis, start_index, narrowed_length)\\n"\n        "    return _align_tensor_to_target_shape(y, target_shape)\\n\\n"\n        "def _apply_strided_slice(x: torch.Tensor, begin: torch.Tensor, end: torch.Tensor, strides: torch.Tensor, begin_mask: int, end_mask: int, target_shape: Optional[Sequence[int]]) -> torch.Tensor:\\n"\n        "    begin_values = _shape_list(begin)\\n"\n        "    end_values = _shape_list(end)\\n"\n        "    stride_values = _shape_list(strides)\\n"\n        "    y = x\\n"\n        "    for axis, (start, stop, step) in enumerate(zip(begin_values, end_values, stride_values)):\\n"\n        "        dim_size = y.shape[axis]\\n"\n        "        has_dynamic_start = (not ((int(begin_mask) >> axis) & 1)) and isinstance(start, torch.Tensor)\\n"\n        "        has_dynamic_stop = (not ((int(end_mask) >> axis) & 1)) and isinstance(stop, torch.Tensor)\\n"\n        "        has_dynamic_step = isinstance(step, torch.Tensor)\\n"\n        "        if (not torch.jit.is_scripting()) and (has_dynamic_start or has_dynamic_stop or has_dynamic_step):\\n"\n        "            dim_size_tensor = torch.scalar_tensor(dim_size, dtype=torch.int64, device=y.device)\\n"\n        "            axis_indices = torch.arange(dim_size, dtype=torch.int64, device=y.device)\\n"\n        "            step_tensor = torch.reshape(torch.as_tensor(step, dtype=torch.int64, device=y.device), [])\\n"\n        "            if ((int(begin_mask) >> axis) & 1):\\n"\n        "                start_tensor = torch.zeros([], dtype=torch.int64, device=y.device)\\n"\n        "            else:\\n"\n        "                start_tensor = torch.reshape(torch.as_tensor(start, dtype=torch.int64, device=y.device), [])\\n"\n        "                start_tensor = torch.where(start_tensor < 0, start_tensor + dim_size_tensor, start_tensor)\\n"\n        "            if ((int(end_mask) >> axis) & 1):\\n"\n        "                stop_tensor = dim_size_tensor\\n"\n        "            else:\\n"\n        "                stop_tensor = torch.reshape(torch.as_tensor(stop, dtype=torch.int64, device=y.device), [])\\n"\n        "                stop_tensor = torch.where(stop_tensor < 0, stop_tensor + dim_size_tensor, stop_tensor)\\n"\n        "            axis_mask = torch.logical_and(axis_indices >= start_tensor, axis_indices < stop_tensor)\\n"\n        "            axis_mask = torch.logical_and(axis_mask, torch.eq(torch.remainder(axis_indices - start_tensor, step_tensor), 0))\\n"\n        "            axis_indices = axis_indices[axis_mask]\\n"\n        "            y = torch.index_select(y, axis, axis_indices)\\n"\n        "            continue\\n"\n        "        step_value = int(step)\\n"\n        "        if step_value == 0:\\n"\n        "            raise RuntimeError(\'STRIDED_SLICE step must be non-zero\')\\n"\n        "        resolved_start = None if ((int(begin_mask) >> axis) & 1) else int(start)\\n"\n        "        resolved_stop = None if ((int(end_mask) >> axis) & 1) else int(stop)\\n"\n        "        if step_value > 0:\\n"\n        "            start_index = 0 if resolved_start is None else int(resolved_start)\\n"\n        "            stop_index = dim_size if resolved_stop is None else int(resolved_stop)\\n"\n        "            if start_index < 0:\\n"\n        "                start_index += dim_size\\n"\n        "            if stop_index < 0:\\n"\n        "                stop_index += dim_size\\n"\n        "        else:\\n"\n        "            start_index = dim_size - 1 if resolved_start is None else int(resolved_start)\\n"\n        "            stop_index = -1 if resolved_stop is None else int(resolved_stop)\\n"\n        "            if start_index < 0:\\n"\n        "                start_index += dim_size\\n"\n        "            if resolved_stop is not None and stop_index < 0:\\n"\n        "                stop_index += dim_size\\n"\n        "        if step_value == 1:\\n"\n        "            narrowed_length = stop_index - start_index\\n"\n        "            y = torch.narrow(y, axis, start_index, narrowed_length)\\n"\n        "            continue\\n"\n        "        axis_indices = torch.arange(start_index, stop_index, step_value, dtype=torch.int64, device=y.device)\\n"\n        "        y = torch.index_select(y, axis, axis_indices)\\n"\n        "    return _align_tensor_to_target_shape(y, target_shape)\\n\\n"\n        "def _resolve_same_padding(input_size: int, kernel_size: int, stride: int, target_size: Optional[int] = None) -> Tuple[int, int]:\\n"\n        "    if target_size is None:\\n"\n        "        out_size = (int(input_size) + int(stride) - 1) // int(stride)\\n"\n        "    else:\\n"\n        "        out_size = int(target_size)\\n"\n        "    total = max((int(out_size) - 1) * int(stride) + int(kernel_size) - int(input_size), 0)\\n"\n        "    before = total // 2\\n"\n        "    after = total - before\\n"\n        "    return before, after\\n\\n"\n        "def _apply_pool2d(x: torch.Tensor, filter_height: int, filter_width: int, stride_h: int, stride_w: int, padding: str, target_shape: Optional[Sequence[int]], is_max_pool: bool, channel_last: Optional[bool] = None) -> torch.Tensor:\\n"\n        "    resize_as_channel_last = bool(channel_last) if channel_last is not None else False\\n"\n        "    if channel_last is None:\\n"\n        "        has_actual_shape, actual_shape = _tensor_static_shape_list(x)\\n"\n        "        if has_actual_shape:\\n"\n        "            resize_as_channel_last = bool(\\n"\n        "                x.ndim == 4\\n"\n        "                and actual_shape[1] > actual_shape[-1]\\n"\n        "                and actual_shape[2] > actual_shape[-1]\\n"\n        "            )\\n"\n        "    has_target_shape, target = _optional_static_shape_list(target_shape)\\n"\n        "    if channel_last is None and x.ndim == 4 and has_target_shape and len(target) == 4:\\n"\n        "        has_actual_shape, actual_shape = _tensor_static_shape_list(x)\\n"\n        "        if has_actual_shape and (\\n"\n        "            (actual_shape[-1] == target[1] and actual_shape[1] != target[1])\\n"\n        "            or (actual_shape[-1] == target[-1] and actual_shape[1] != target[-1])\\n"\n        "        ):\\n"\n        "            resize_as_channel_last = True\\n"\n        "    pool_input = x.permute(0, 3, 1, 2).contiguous() if resize_as_channel_last and x.ndim == 4 else x\\n"\n        "    if str(padding).upper() == \'SAME\':\\n"\n        "        target_h: Optional[int] = None\\n"\n        "        target_w: Optional[int] = None\\n"\n        "        if has_target_shape and len(target) == 4:\\n"\n        "            if resize_as_channel_last:\\n"\n        "                if int(target[1]) == int(pool_input.shape[1]) and int(target[-1]) != int(pool_input.shape[1]):\\n"\n        "                    target_h = int(target[2])\\n"\n        "                    target_w = int(target[3])\\n"\n        "                else:\\n"\n        "                    target_h = int(target[1])\\n"\n        "                    target_w = int(target[2])\\n"\n        "            else:\\n"\n        "                target_h = int(target[2])\\n"\n        "                target_w = int(target[3])\\n"\n        "        pad_w = _resolve_same_padding(int(pool_input.shape[-1]), filter_width, stride_w, target_w)\\n"\n        "        pad_h = _resolve_same_padding(int(pool_input.shape[-2]), filter_height, stride_h, target_h)\\n"\n        "        pool_input = F.pad(pool_input, [pad_w[0], pad_w[1], pad_h[0], pad_h[1]], mode=\'constant\', value=float(\'-inf\') if is_max_pool else 0.0)\\n"\n        "        padding_value = 0\\n"\n        "    else:\\n"\n        "        padding_value = 0\\n"\n        "    if is_max_pool:\\n"\n        "        y = F.max_pool2d(pool_input, kernel_size=(filter_height, filter_width), stride=(stride_h, stride_w), padding=padding_value)\\n"\n        "    else:\\n"\n        "        y = F.avg_pool2d(pool_input, kernel_size=(filter_height, filter_width), stride=(stride_h, stride_w), padding=padding_value)\\n"\n        "    if resize_as_channel_last and y.ndim == 4:\\n"\n        "        y = y.permute(0, 2, 3, 1).contiguous()\\n"\n        "    return _align_tensor_to_target_shape(y, target_shape)\\n\\n"\n        "@torch.jit.ignore\\n"\n        "def _onnx_export_active() -> bool:\\n"\n        "    return bool(torch.onnx.is_in_onnx_export())\\n\\n"\n        "def _resize_bilinear_exact(x: torch.Tensor, size: Sequence[int], *, align_corners: bool, half_pixel_centers: bool) -> torch.Tensor:\\n"\n        "    if not torch.jit.is_scripting():\\n"\n        "        if _onnx_export_active():\\n"\n        "            return F.interpolate(x, size=[int(size[0]), int(size[1])], mode=\'bilinear\', align_corners=align_corners)\\n"\n        "    if x.ndim != 4:\\n"\n        "        return F.interpolate(x, size=[int(size[0]), int(size[1])], mode=\'bilinear\', align_corners=align_corners)\\n"\n        "    out_h = int(size[0])\\n"\n        "    out_w = int(size[1])\\n"\n        "    in_h = int(x.shape[-2])\\n"\n        "    in_w = int(x.shape[-1])\\n"\n        "    if out_h <= 0 or out_w <= 0:\\n"\n        "        raise RuntimeError(\'Resize target dimensions must be positive.\')\\n"\n        "    if align_corners:\\n"\n        "        ys = torch.zeros([out_h], dtype=torch.float32, device=x.device) if out_h == 1 else torch.arange(out_h, dtype=torch.float32, device=x.device) * float(max(in_h - 1, 0)) / float(max(out_h - 1, 1))\\n"\n        "        xs = torch.zeros([out_w], dtype=torch.float32, device=x.device) if out_w == 1 else torch.arange(out_w, dtype=torch.float32, device=x.device) * float(max(in_w - 1, 0)) / float(max(out_w - 1, 1))\\n"\n        "    elif half_pixel_centers:\\n"\n        "        ys = (torch.arange(out_h, dtype=torch.float32, device=x.device) + 0.5) * float(in_h) / float(out_h) - 0.5\\n"\n        "        xs = (torch.arange(out_w, dtype=torch.float32, device=x.device) + 0.5) * float(in_w) / float(out_w) - 0.5\\n"\n        "    else:\\n"\n        "        ys = torch.arange(out_h, dtype=torch.float32, device=x.device) * float(in_h) / float(out_h)\\n"\n        "        xs = torch.arange(out_w, dtype=torch.float32, device=x.device) * float(in_w) / float(out_w)\\n"\n        "    y0 = torch.floor(ys).to(dtype=torch.int64)\\n"\n        "    x0 = torch.floor(xs).to(dtype=torch.int64)\\n"\n        "    y1 = y0 + 1\\n"\n        "    x1 = x0 + 1\\n"\n        "    y0c = y0.clamp(0, max(in_h - 1, 0))\\n"\n        "    x0c = x0.clamp(0, max(in_w - 1, 0))\\n"\n        "    y1c = y1.clamp(0, max(in_h - 1, 0))\\n"\n        "    x1c = x1.clamp(0, max(in_w - 1, 0))\\n"\n        "    ly = (ys - y0.to(dtype=torch.float32)).view(1, 1, out_h, 1)\\n"\n        "    lx = (xs - x0.to(dtype=torch.float32)).view(1, 1, 1, out_w)\\n"\n        "    hy = 1.0 - ly\\n"\n        "    hx = 1.0 - lx\\n"\n        "    top_left = x[:, :, y0c[:, None], x0c[None, :]]\\n"\n        "    top_right = x[:, :, y0c[:, None], x1c[None, :]]\\n"\n        "    bottom_left = x[:, :, y1c[:, None], x0c[None, :]]\\n"\n        "    bottom_right = x[:, :, y1c[:, None], x1c[None, :]]\\n"\n        "    return top_left * hy * hx + top_right * hy * lx + bottom_left * ly * hx + bottom_right * ly * lx\\n\\n"\n        "def _apply_resize(x: torch.Tensor, size: Any, method: str, target_shape: Optional[Sequence[int]], align_corners: bool = False, half_pixel_centers: bool = False, channel_last: Optional[bool] = None) -> torch.Tensor:\\n"\n        "    resize_size = _shape_list(size)\\n"\n        "    resize_as_channel_last = bool(channel_last) if channel_last is not None else False\\n"\n        "    has_target_shape, target = _optional_static_shape_list(target_shape)\\n"\n        "    if channel_last is None and x.ndim == 4 and has_target_shape and len(target) == 4:\\n"\n        "        has_actual_shape, actual_shape = _tensor_static_shape_list(x)\\n"\n        "        if has_actual_shape and (\\n"\n        "            (actual_shape[-1] == target[1] and actual_shape[1] != target[1])\\n"\n        "            or (actual_shape[-1] == target[-1] and actual_shape[1] != target[-1])\\n"\n        "        ):\\n"\n        "            resize_as_channel_last = True\\n"\n        "    resize_input = x.permute(0, 3, 1, 2).contiguous() if resize_as_channel_last and x.ndim == 4 else x\\n"\n        "    if str(method).lower() == \'nearest\':\\n"\n        "        y = F.interpolate(resize_input, size=resize_size, mode=\'nearest\')\\n"\n        "    else:\\n"\n        "        y = _resize_bilinear_exact(resize_input, resize_size, align_corners=align_corners, half_pixel_centers=half_pixel_centers)\\n"\n        "    if resize_as_channel_last and y.ndim == 4:\\n"\n        "        y = y.permute(0, 2, 3, 1).contiguous()\\n"\n        "    return _align_tensor_to_target_shape(y, target_shape)\\n\\n"\n        "def _box_iou(boxes: torch.Tensor, box: torch.Tensor) -> torch.Tensor:\\n"\n        "    x1 = torch.maximum(boxes[:, 0], box[0])\\n"\n        "    y1 = torch.maximum(boxes[:, 1], box[1])\\n"\n        "    x2 = torch.minimum(boxes[:, 2], box[2])\\n"\n        "    y2 = torch.minimum(boxes[:, 3], box[3])\\n"\n        "    inter_w = torch.clamp(x2 - x1, min=0.0)\\n"\n        "    inter_h = torch.clamp(y2 - y1, min=0.0)\\n"\n        "    inter = inter_w * inter_h\\n"\n        "    boxes_area = torch.clamp(boxes[:, 2] - boxes[:, 0], min=0.0) * torch.clamp(boxes[:, 3] - boxes[:, 1], min=0.0)\\n"\n        "    box_area = torch.clamp(box[2] - box[0], min=0.0) * torch.clamp(box[3] - box[1], min=0.0)\\n"\n        "    union = boxes_area + box_area - inter\\n"\n        "    safe_union = torch.where(union > 0, union, torch.ones_like(union))\\n"\n        "    iou = inter / safe_union\\n"\n        "    return torch.where(union > 0, iou, torch.zeros_like(iou))\\n\\n"\n        "def _apply_non_max_suppression_v4(boxes: torch.Tensor, scores: torch.Tensor, max_output_size: torch.Tensor, iou_threshold: torch.Tensor, score_threshold: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:\\n"\n        "    flat_boxes = boxes.to(dtype=torch.float32).reshape(-1, 4)\\n"\n        "    flat_scores = scores.to(dtype=torch.float32).reshape(-1)\\n"\n        "    max_outputs = max(0, int(max_output_size.reshape(-1)[0].to(dtype=torch.int64).item()))\\n"\n        "    selected = torch.zeros([max_outputs], dtype=torch.int64, device=flat_boxes.device)\\n"\n        "    valid_count = torch.zeros([], dtype=torch.int32, device=flat_boxes.device)\\n"\n        "    if max_outputs == 0:\\n"\n        "        return selected.to(dtype=torch.int32), valid_count\\n"\n        "    iou_thresh = float(iou_threshold.reshape(-1)[0].item())\\n"\n        "    score_thresh = float(score_threshold.reshape(-1)[0].item())\\n"\n        "    candidate_scores = torch.where(\\n"\n        "        flat_scores > score_thresh,\\n"\n        "        flat_scores,\\n"\n        "        torch.full_like(flat_scores, float(\'-inf\')),\\n"\n        "    )\\n"\n        "    all_indices = torch.arange(flat_scores.shape[0], dtype=torch.int64, device=flat_boxes.device)\\n"\n        "    neg_inf = torch.full_like(candidate_scores, float(\'-inf\'))\\n"\n        "    for out_index in range(max_outputs):\\n"\n        "        current_score, current_index = torch.max(candidate_scores, dim=0)\\n"\n        "        current_index = current_index.to(dtype=torch.int64)\\n"\n        "        is_valid = torch.isfinite(current_score)\\n"\n        "        selected[out_index : out_index + 1] = torch.where(is_valid, current_index, torch.zeros_like(current_index)).reshape(1)\\n"\n        "        valid_count = valid_count + is_valid.to(dtype=torch.int32)\\n"\n        "        current_box = torch.index_select(flat_boxes, 0, current_index.reshape(1)).reshape(4)\\n"\n        "        suppress = _box_iou(flat_boxes, current_box) > iou_thresh\\n"\n        "        suppress = torch.logical_or(suppress, all_indices == current_index)\\n"\n        "        suppress = torch.logical_and(suppress, is_valid)\\n"\n        "        candidate_scores = torch.where(suppress, neg_inf, candidate_scores)\\n"\n        "    return selected.to(dtype=torch.int32), valid_count\\n\\n"\n        "def _crop_nms_selected_indices(selected_indices: torch.Tensor, valid_count: torch.Tensor) -> torch.Tensor:\\n"\n        "    limit = valid_count.to(dtype=torch.int64).reshape([])\\n"\n        "    return torch.narrow(selected_indices, 0, 0, limit)\\n\\n"\n        "def _normalize_axes(value: Any, rank: int) -> Optional[List[int]]:\\n"\n        "    if value is None:\\n"\n        "        return None\\n"\n        "    axes = _shape_list(value)\\n"\n        "    normalized = torch.jit.annotate(List[int], [])\\n"\n        "    for raw_axis in axes:\\n"\n        "        axis = _normalize_dim(int(raw_axis), rank)\\n"\n        "        insert_at = 0\\n"\n        "        while insert_at < len(normalized) and int(normalized[insert_at]) < axis:\\n"\n        "            insert_at += 1\\n"\n        "        if insert_at < len(normalized) and int(normalized[insert_at]) == axis:\\n"\n        "            continue\\n"\n        "        normalized.insert(insert_at, axis)\\n"\n        "    return normalized\\n\\n"\n        "def _reduce_sum(x: torch.Tensor, axis: Optional[List[int]], keepdims: bool) -> torch.Tensor:\\n"\n        "    if axis is None:\\n"\n        "        return torch.sum(x) if not keepdims else torch.sum(x).reshape([1] * x.ndim)\\n"\n        "    result = x\\n"\n        "    if keepdims:\\n"\n        "        for dim in axis:\\n"\n        "            result = torch.sum(result, dim=int(dim), keepdim=True)\\n"\n        "        return result\\n"\n        "    reverse_index = len(axis) - 1\\n"\n        "    while reverse_index >= 0:\\n"\n        "        result = torch.sum(result, dim=int(axis[reverse_index]), keepdim=False)\\n"\n        "        reverse_index -= 1\\n"\n        "    return result\\n\\n"\n        "def _reduce_mean(x: torch.Tensor, axis: Optional[List[int]], keepdims: bool) -> torch.Tensor:\\n"\n        "    if axis is None:\\n"\n        "        return torch.mean(x) if not keepdims else torch.mean(x).reshape([1] * x.ndim)\\n"\n        "    result = x\\n"\n        "    if keepdims:\\n"\n        "        for dim in axis:\\n"\n        "            result = torch.mean(result, dim=int(dim), keepdim=True)\\n"\n        "        return result\\n"\n        "    reverse_index = len(axis) - 1\\n"\n        "    while reverse_index >= 0:\\n"\n        "        result = torch.mean(result, dim=int(axis[reverse_index]), keepdim=False)\\n"\n        "        reverse_index -= 1\\n"\n        "    return result\\n\\n"\n        "def _reduce_max(x: torch.Tensor, axis: Optional[List[int]], keepdims: bool) -> torch.Tensor:\\n"\n        "    if axis is None:\\n"\n        "        return torch.amax(x, keepdim=keepdims)\\n"\n        "    result = x\\n"\n        "    if keepdims:\\n"\n        "        for dim in axis:\\n"\n        "            result = torch.amax(result, dim=int(dim), keepdim=True)\\n"\n        "        return result\\n"\n        "    reverse_index = len(axis) - 1\\n"\n        "    while reverse_index >= 0:\\n"\n        "        result = torch.amax(result, dim=int(axis[reverse_index]), keepdim=False)\\n"\n        "        reverse_index -= 1\\n"\n        "    return result\\n\\n"\n        "def _reduce_min(x: torch.Tensor, axis: Optional[List[int]], keepdims: bool) -> torch.Tensor:\\n"\n        "    if axis is None:\\n"\n        "        return torch.amin(x, keepdim=keepdims)\\n"\n        "    result = x\\n"\n        "    if keepdims:\\n"\n        "        for dim in axis:\\n"\n        "            result = torch.amin(result, dim=int(dim), keepdim=True)\\n"\n        "        return result\\n"\n        "    reverse_index = len(axis) - 1\\n"\n        "    while reverse_index >= 0:\\n"\n        "        result = torch.amin(result, dim=int(axis[reverse_index]), keepdim=False)\\n"\n        "        reverse_index -= 1\\n"\n        "    return result\\n\\n"\n        "def _reduce_prod(x: torch.Tensor, axis: Optional[List[int]], keepdims: bool) -> torch.Tensor:\\n"\n        "    if axis is None:\\n"\n        "        y = torch.prod(x)\\n"\n        "        return y if not keepdims else y.reshape([1] * x.ndim)\\n"\n        "    result = x\\n"\n        "    if keepdims:\\n"\n        "        for dim in axis:\\n"\n        "            result = torch.prod(result, dim=int(dim), keepdim=True)\\n"\n        "        return result\\n"\n        "    reverse_index = len(axis) - 1\\n"\n        "    while reverse_index >= 0:\\n"\n        "        result = torch.prod(result, dim=int(axis[reverse_index]), keepdim=False)\\n"\n        "        reverse_index -= 1\\n"\n        "    return result\\n\\n"\n        "def _reduce_any(x: torch.Tensor, axis: Optional[List[int]], keepdims: bool) -> torch.Tensor:\\n"\n        "    if axis is None:\\n"\n        "        y = torch.any(x)\\n"\n        "        return y if not keepdims else y.reshape([1] * x.ndim)\\n"\n        "    result = x\\n"\n        "    if keepdims:\\n"\n        "        for dim in axis:\\n"\n        "            result = torch.any(result, dim=int(dim), keepdim=True)\\n"\n        "        return result\\n"\n        "    reverse_index = len(axis) - 1\\n"\n        "    while reverse_index >= 0:\\n"\n        "        result = torch.any(result, dim=int(axis[reverse_index]), keepdim=False)\\n"\n        "        reverse_index -= 1\\n"\n        "    return result\\n\\n"\n    ) + conv_block_helper_source\n\n    runtime_source = _build_native_runtime_source(helper_source)\n    _write_generated_package_common_files(\n        output_folder_path,\n        runtime_source=runtime_source,\n    )\n    buffer_init_block = "\\n".join(f"        {line}" for line in buffer_init_lines)\n    buffer_annotation_block = ""\n    init_constants_method = ""\n    if len(buffer_init_lines) > 0:\n        init_constants_method = (\n            "    def _init_constants(self) -> None:\\n"\n            f"{buffer_init_block}\\n\\n"\n        )\n    init_constants_call = "        self._init_constants()\\n" if len(buffer_init_lines) > 0 else ""\n    forward_kwargs_lines: List[str] = []\n    forward_args_lines: List[str] = []\n    forward_signature_args: List[str] = []\n    for input_index, input_name in enumerate(model_ir.inputs):\n        input_var = tensor_var_names[str(input_name)]\n        forward_signature_args.append(f"{input_var}: torch.Tensor")\n        forward_kwargs_lines.append(\n            f"            {input_var} = resolve_named_input_value(kwargs, {str(input_name)!r})"\n        )\n        forward_args_lines.append(\n            f"            {input_var} = args[{input_index}]"\n        )\n\n    forward_lines = _fold_channel_last_affine_conv_bridges(\n        forward_lines,\n        derive_local_var_name=_derived_local_var_name,\n        channel_first_constant_expr_for_buffer_attr=_channel_first_constant_expr_for_buffer_attr,\n    )\n    forward_lines = _rewrite_channel_last_binary_bridge_chains(\n        forward_lines,\n        derive_local_var_name=_derived_local_var_name,\n        channel_first_constant_expr_for_buffer_attr=_channel_first_constant_expr_for_buffer_attr,\n    )\n    forward_lines = _fold_channel_first_gap_conv_bridges(forward_lines)\n    forward_lines = _rewrite_channel_first_gap_outputs_to_explicit_channel_last(forward_lines)\n    forward_lines = _rewrite_channel_last_gap_means_to_reduce_mean(forward_lines)\n    forward_lines = _fold_boundary_transpose_pad_conv_bridges(forward_lines)\n    bridged_forward_lines = _bridge_boundary_metadata_gather_nd_inputs(\n        forward_lines,\n        model_ir=model_ir,\n        tensor_var_names=tensor_var_names,\n    )\n    if bridged_forward_lines != forward_lines:\n        runtime_imports.add("_torch_permute")\n    forward_lines = bridged_forward_lines\n    collapsed_forward_lines = _collapse_redundant_torch_permute_chains(forward_lines)\n    if collapsed_forward_lines != forward_lines:\n        runtime_imports.add("_torch_permute")\n    forward_lines = collapsed_forward_lines\n    forward_lines = _fold_channel_last_prelu_bridges(forward_lines)\n    forward_lines = _inline_trivial_public_layout_bridge_aliases(forward_lines)\n    forward_lines = _fold_rank4_reshape_permute_conv_bridges(forward_lines)\n    forward_lines = _fold_single_use_static_reshape_chains(\n        forward_lines,\n        tensor_var_names=tensor_var_names,\n        model_ir=model_ir,\n    )\n    forward_lines = _prune_dead_forward_lines(\n        forward_lines,\n        input_var_names=[str(tensor_var_names[str(name)]) for name in model_ir.inputs],\n        output_var_names=[str(tensor_var_names[str(name)]) for name in model_ir.outputs],\n    )\n    stage_methods_source, forward_stage_calls, stage_specs = _build_forward_stage_methods(\n        forward_lines,\n        tensor_var_names=tensor_var_names,\n        model_ir=model_ir,\n    )\n    named_encoder_class_source = ""\n    named_encoder_init_lines: List[str] = []\n    if len(stage_specs) > 0:\n        named_encoder_class_source, named_encoder_init_lines, forward_stage_calls = _build_named_encoder_methods_composite(\n            stage_specs,\n            final_output_names={str(tensor_var_names[str(name)]) for name in model_ir.outputs},\n        )\n    bridged_forward_stage_calls = _bridge_boundary_metadata_gather_nd_inputs(\n        [str(line).strip() for line in forward_stage_calls],\n        model_ir=model_ir,\n        tensor_var_names=tensor_var_names,\n    )\n    collapsed_forward_stage_calls = _collapse_redundant_torch_permute_chains(\n        bridged_forward_stage_calls\n    )\n    if collapsed_forward_stage_calls != bridged_forward_stage_calls:\n        bridged_forward_stage_calls = collapsed_forward_stage_calls\n    bridged_forward_stage_calls = _inline_trivial_public_layout_bridge_aliases(\n        bridged_forward_stage_calls\n    )\n    if bridged_forward_stage_calls != [str(line).strip() for line in forward_stage_calls]:\n        runtime_imports.add("_torch_permute")\n        forward_stage_calls = [f"        {line}" for line in bridged_forward_stage_calls]\n    if len(named_encoder_init_lines) > 0:\n        module_init_lines.extend(named_encoder_init_lines)\n    module_init_block = "\\n".join(f"        {line}" for line in module_init_lines)\n    forward_block = "\\n".join(forward_stage_calls)\n    forward_kwargs_block = "\\n".join(forward_kwargs_lines) if len(forward_kwargs_lines) > 0 else "            pass"\n    forward_args_block = "\\n".join(forward_args_lines) if len(forward_args_lines) > 0 else "            pass"\n    forward_signature = ", ".join(["self"] + forward_signature_args)\n    forward_named_call_args = ", ".join(\n        str(tensor_var_names[str(input_name)])\n        for input_name in model_ir.inputs\n    )\n    outputs_expr = ", ".join(\n        (\n            tensor_var_names[str(name)]\n            if str(name) in producer_index or str(name) in model_ir.inputs\n            else _tensor_expr(str(name))\n        )\n        for name in model_ir.outputs\n    )\n    runtime_import_order = [\n        "_Conv2dBlock",\n        "_align_binary_inputs",\n        "_align_binary_inputs_to_anchor",\n        "_align_tensor_to_target_shape",\n        "_apply_concat",\n        "_apply_cumsum",\n        "_apply_if_axis0_tensor_mux",\n        "_crop_nms_selected_indices",\n        "_apply_fused_activation",\n        "_apply_gather",\n        "_apply_gather_elements",\n        "_apply_gather_nd",\n        "_apply_random_standard_normal",\n        "_apply_scatter_nd",\n        "_apply_tile",\n        "_apply_module_conv2d",\n        "_apply_module_conv3d",\n        "_apply_module_transpose_conv2d",\n        "_apply_module_transpose_conv3d",\n        "_apply_non_max_suppression_v4",\n        "_apply_pool2d",\n        "_apply_pad_nd",\n        "_apply_resize",\n        "_apply_slice",\n        "_apply_softmax",\n        "_apply_strided_slice",\n        "_coerce_scalar_axis",\n        "_module_device",\n        "_normalize_axes",\n        "_normalize_dim",\n        "_reduce_any",\n        "_reduce_max",\n        "_reduce_mean",\n        "_reduce_min",\n        "_reduce_prod",\n        "_reduce_sum",\n        "_reshape_gather_output",\n        "_resolve_reshape_shape",\n        "_resolve_reshape_shape_tensor",\n        "resolve_named_input_value",\n        "_shape_list",\n        "_tensor_shape_list",\n        "_shape_tensor",\n        "_to_torch_pad_arg",\n        "_torch_dtype",\n        "_torch_permute",\n        "load_generated_weights",\n    ]\n    if has_conv_blocks:\n        runtime_imports.add("_Conv2dBlock")\n    has_affine_layer_norms = len(affine_layer_norm_specs) > 0\n    affine_layer_norm_source = (\n        "class _AffineLayerNorm(torch.nn.Module):\\n"\n        "    gamma: torch.Tensor\\n"\n        "    beta: torch.Tensor\\n\\n"\n        "    def __init__(self, *, shape: list[int], dtype: torch.dtype) -> None:\\n"\n        "        super().__init__()\\n"\n        "        self.register_buffer(\'gamma\', torch.zeros(shape, dtype=dtype), persistent=True)\\n"\n        "        self.register_buffer(\'beta\', torch.zeros(shape, dtype=dtype), persistent=True)\\n\\n"\n        "    def forward(self, x: torch.Tensor) -> torch.Tensor:\\n"\n        "        return torch.add(torch.mul(x, self.gamma), self.beta)\\n\\n"\n    ) if has_affine_layer_norms else ""\n    nms_method_source = ""\n    if len(nms_method_specs) > 0:\n        runtime_imports.add("_module_device")\n        method_chunks: List[str] = []\n        for spec in nms_method_specs:\n            extra_args: List[str] = []\n            if spec.get("max_output_arg_expr", None) is not None:\n                extra_args.append("max_output_size: torch.Tensor")\n            if spec.get("iou_threshold_arg_expr", None) is not None:\n                extra_args.append("iou_threshold: torch.Tensor")\n            if spec.get("score_threshold_arg_expr", None) is not None:\n                extra_args.append("score_threshold: torch.Tensor")\n            extra_signature = ""\n            if len(extra_args) > 0:\n                extra_signature = ", " + ", ".join(extra_args)\n            method_chunks.append(\n                "    def {name}(self, boxes: torch.Tensor, scores: torch.Tensor{extra_signature}) -> tuple[torch.Tensor, torch.Tensor]:\\n"\n                "        return _apply_non_max_suppression_v4(\\n"\n                "            boxes,\\n"\n                "            scores,\\n"\n                "            torch.as_tensor({max_output_expr}, dtype=torch.int32, device=_module_device(self)),\\n"\n                "            torch.as_tensor({iou_threshold_expr}, dtype=torch.float32, device=_module_device(self)),\\n"\n                "            torch.as_tensor({score_threshold_expr}, dtype=torch.float32, device=_module_device(self)),\\n"\n                "        )\\n\\n".format(\n                    name=str(spec["name"]),\n                    extra_signature=extra_signature,\n                    max_output_expr=str(spec["max_output_expr"]),\n                    iou_threshold_expr=str(spec["iou_threshold_expr"]),\n                    score_threshold_expr=str(spec["score_threshold_expr"]),\n                )\n            )\n        nms_method_source = "".join(method_chunks)\n    runtime_import_block = "".join(\n        f"    {name},\\n" for name in runtime_import_order if name in runtime_imports\n    )\n    constant_buffer_alias_method_source = ""\n    constant_buffer_alias_init_call = ""\n    constant_buffer_alias_load_state_dict_source = ""\n    if (\n        len(channel_first_constant_buffer_alias_refresh_specs) > 0\n        or len(permuted_constant_buffer_alias_refresh_specs) > 0\n        or len(transposed_constant_buffer_alias_refresh_specs) > 0\n    ):\n        refresh_lines = [\n            "    def _refresh_constant_buffer_aliases(self) -> None:\\n",\n            "        with torch.no_grad():\\n",\n        ]\n        for alias_attr_name, source_attr_name, alias_shape in channel_first_constant_buffer_alias_refresh_specs:\n            refresh_lines.append(\n                f"            self.{alias_attr_name}.copy_(torch.reshape(self.{source_attr_name}, {repr(alias_shape)}))\\n"\n            )\n        for alias_attr_name, source_attr_name, perm in permuted_constant_buffer_alias_refresh_specs:\n            refresh_lines.append(\n                f"            self.{alias_attr_name}.copy_(self.{source_attr_name}.permute(*{repr(tuple(int(v) for v in list(perm)))}).contiguous())\\n"\n            )\n        for alias_attr_name, source_attr_name in transposed_constant_buffer_alias_refresh_specs:\n            refresh_lines.append(\n                f"            self.{alias_attr_name}.copy_(self.{source_attr_name}.transpose(-1, -2))\\n"\n            )\n        refresh_lines.append("\\n")\n        constant_buffer_alias_method_source = "".join(refresh_lines)\n        constant_buffer_alias_init_call = "        self._refresh_constant_buffer_aliases()\\n"\n        constant_buffer_alias_load_state_dict_source = (\n            "    def load_state_dict(self, state_dict: Dict[str, torch.Tensor], strict: bool = True, assign: bool = False):\\n"\n            "        result = super().load_state_dict(state_dict, strict=strict, assign=assign)\\n"\n            "        self._refresh_constant_buffer_aliases()\\n"\n            "        return result\\n\\n"\n        )\n\n    model_source = (\n        "# pyright: reportArgumentType=false, reportCallIssue=false\\n"\n        "from pathlib import Path\\n"\n        "from typing import Any, Callable, Dict, Optional, Tuple\\n\\n"\n        "import torch\\n"\n        "import torch.nn.functional as F\\n\\n"\n        "from .runtime import (\\n"\n        f"{runtime_import_block}"\n        ")\\n\\n"\n        "PACKAGE_DIR = Path(__file__).resolve().parent\\n"\n        f"INPUT_NAMES = {repr([str(v) for v in model_ir.inputs])}\\n"\n        f"OUTPUT_NAMES = {repr([str(v) for v in model_ir.outputs])}\\n"\n        f"{sequence_rnn_helper_source}"\n        f"{sequence_lstm_helper_source}"\n        f"{affine_layer_norm_source}"\n        f"{named_encoder_class_source}"\n        "class Model(torch.nn.Module):\\n"\n        f"{buffer_annotation_block}"\n        "    def __init__(self, *, device: str | None = None, eval_mode: bool = True, load_weights: bool = True):\\n"\n        "        super().__init__()\\n"\n        "        self.input_names = list(INPUT_NAMES)\\n"\n        "        self.output_names = list(OUTPUT_NAMES)\\n"\n        "        self._onnx2tf_torch_export_mode = False\\n"\n        f"{module_init_block}\\n"\n        f"{init_constants_call}"\n        "        if load_weights:\\n"\n        "            load_generated_weights(\\n"\n        "                model=self,\\n"\n        "                package_dir=PACKAGE_DIR,\\n"\n        "                device=device,\\n"\n        "            )\\n"\n        "        elif device is not None:\\n"\n        "            self.to(device)\\n"\n        f"{constant_buffer_alias_init_call}"\n        "        if eval_mode:\\n"\n        "            self.eval()\\n\\n"\n        f"{init_constants_method}"\n        f"{constant_buffer_alias_method_source}"\n        f"{constant_buffer_alias_load_state_dict_source}"\n        f"{nms_method_source}"\n        f"{stage_methods_source}"\n        f"    def forward({forward_signature}) -> Any:\\n"\n        f"{forward_block}\\n"\n    )\n    if len(model_ir.outputs) == 1:\n        model_source += (\n            f"        return {outputs_expr}\\n\\n"\n            "    def forward_named(self, *args: torch.Tensor, **kwargs: torch.Tensor) -> Dict[str, torch.Tensor]:\\n"\n            "        if len(args) > 0 and len(kwargs) > 0:\\n"\n            "            raise RuntimeError(\'Use either positional inputs or keyword inputs, not both.\')\\n"\n            "        if len(kwargs) > 0:\\n"\n            f"{forward_kwargs_block}\\n"\n            "        else:\\n"\n            f"            if len(args) != {len(model_ir.inputs)}:\\n"\n            "                raise RuntimeError(f\'Input arity mismatch. expected={len(self.input_names)} actual={len(args)}\')\\n"\n            f"{forward_args_block}\\n"\n            f"        return {{{str(model_ir.outputs[0])!r}: self.forward({forward_named_call_args})}}\\n\\n"\n        )\n    else:\n        model_source += (\n            f"        return ({outputs_expr})\\n\\n"\n            "    def forward_named(self, *args: torch.Tensor, **kwargs: torch.Tensor) -> Dict[str, torch.Tensor]:\\n"\n            "        if len(args) > 0 and len(kwargs) > 0:\\n"\n            "            raise RuntimeError(\'Use either positional inputs or keyword inputs, not both.\')\\n"\n            "        if len(kwargs) > 0:\\n"\n            f"{forward_kwargs_block}\\n"\n            "        else:\\n"\n            f"            if len(args) != {len(model_ir.inputs)}:\\n"\n            "                raise RuntimeError(f\'Input arity mismatch. expected={len(self.input_names)} actual={len(args)}\')\\n"\n            f"{forward_args_block}\\n"\n            f"        result = self.forward({forward_named_call_args})\\n"\n            f"        return {{name: value for name, value in zip({repr([str(v) for v in model_ir.outputs])}, result)}}\\n\\n"\n        )\n    model_source += (\n        "def load_model(device: str | None = None, eval_mode: bool = True) -> Model:\\n"\n        "    return Model(device=device, eval_mode=eval_mode)\\n"\n    )\n    (package_dir / "model.py").write_text(model_source, encoding="utf-8")\n    return [(str(attr_path), str(tensor_name)) for attr_path, tensor_name in load_specs]\n\n'
_NATIVE_CODEGEN_FUNCTION_SOURCE = _NATIVE_CODEGEN_FUNCTION_SOURCE.replace(
    "state_dict: Dict[str, torch.Tensor]",
    "state_dict: Mapping[str, Any]",
).replace(
    "from typing import Any, Callable, Dict, Optional, Tuple",
    "from typing import Any, Callable, Dict, Mapping, Optional, Tuple",
).replace(
    '''resize_channel_last_expr = (
                "False"
                if resize_uses_channel_first_alias else
                "True" if resize_shape_implies_channel_last is True else
                "False" if resize_shape_implies_channel_last is False else
                "True" if resize_effective_layout == "NHWC" else
                "False" if resize_effective_layout == "NCHW" else
                "None"
            )''',
    '''resize_channel_last_expr = (
                "False"
                if resize_uses_channel_first_alias else
                "True" if resize_effective_layout == "NHWC" else
                "False" if resize_effective_layout == "NCHW" else
                "True" if resize_shape_implies_channel_last is True else
                "False" if resize_shape_implies_channel_last is False else
                "None"
            )''',
).replace(
    "if can_emit_channel_first_binary_op_fn(op) and runtime_shape_passthrough_operand is None and not requires_runtime_alignment:",
    "if can_emit_channel_first_binary_op_fn(op) and runtime_shape_passthrough_operand is None:",
)
_NATIVE_CODEGEN_IMPL: Optional[Callable[..., Any]] = None


def _rewrite_public_layout_bridge_binary_align_calls(model_file: Path) -> None:
    if not model_file.exists():
        return
    source = model_file.read_text(encoding="utf-8")
    bridge_assignments: Dict[str, list[int]] = {}
    for match in re.finditer(
        r"(?m)^\s*(?P<bridge>[A-Za-z0-9_]*layout_bridge)\s*=\s*_torch_permute\((?P<source>[A-Za-z0-9_]+), \[(?P<perm>[0-9,\s]+)\]\)\s*$",
        source,
    ):
        perm = [int(token.strip()) for token in str(match.group("perm")).split(",")]
        if sorted(perm) != list(range(len(perm))):
            continue
        inverse_perm = [0 for _ in perm]
        for axis, value in enumerate(perm):
            inverse_perm[int(value)] = int(axis)
        bridge_assignments[str(match.group("bridge"))] = inverse_perm
    if len(bridge_assignments) == 0:
        return
    rewritten = source
    for bridge_name, inverse_perm in bridge_assignments.items():
        bridge_pattern = re.escape(bridge_name)
        perm_literal = "[" + ", ".join(str(int(v)) for v in inverse_perm) + "]"

        def _wrap_other_arg(match: re.Match[str]) -> str:
            other_expr = str(match.group("other")).strip()
            if other_expr.startswith("_torch_permute("):
                return str(match.group(0))
            return (
                f"_align_binary_inputs(_torch_permute({other_expr}, {perm_literal}), "
                f"{bridge_name}, {match.group('target')})"
            )

        def _wrap_other_arg_rhs(match: re.Match[str]) -> str:
            other_expr = str(match.group("other")).strip()
            if other_expr.startswith("_torch_permute("):
                return str(match.group(0))
            return (
                f"_align_binary_inputs({bridge_name}, "
                f"_torch_permute({other_expr}, {perm_literal}), {match.group('target')})"
            )

        rewritten = re.sub(
            rf"_align_binary_inputs\((?P<other>[^,\n]+),\s*{bridge_pattern},\s*(?P<target>\[[^\]]+\])\)",
            _wrap_other_arg,
            rewritten,
        )
        rewritten = re.sub(
            rf"_align_binary_inputs\({bridge_pattern},\s*(?P<other>[^,\n]+),\s*(?P<target>\[[^\]]+\])\)",
            _wrap_other_arg_rhs,
            rewritten,
        )
    if rewritten != source:
        model_file.write_text(rewritten, encoding="utf-8")


def _resolve_static_split_axis(
    input_shape: list[int],
    output_shapes: list[list[int]],
) -> Optional[int]:
    rank = len(input_shape)
    if rank == 0 or len(output_shapes) == 0:
        return None
    candidate_axes: list[int] = []
    for axis in range(rank):
        valid = True
        axis_extent_static = int(input_shape[axis]) > 0
        expected_axis_extent = 0
        for output_shape in output_shapes:
            if len(output_shape) != rank:
                valid = False
                break
            for dim_index, (input_dim, output_dim) in enumerate(zip(input_shape, output_shape)):
                if dim_index == axis:
                    continue
                if int(input_dim) > 0 and int(output_dim) > 0 and int(input_dim) != int(output_dim):
                    valid = False
                    break
            if not valid:
                break
            if int(output_shape[axis]) <= 0:
                axis_extent_static = False
            else:
                expected_axis_extent += int(output_shape[axis])
        if valid and (not axis_extent_static or expected_axis_extent == int(input_shape[axis])):
            candidate_axes.append(int(axis))
    if len(candidate_axes) == 1:
        return int(candidate_axes[0])
    return None


def _rewrite_split_axes_from_static_shapes(
    model_file: Path,
    *,
    model_ir: Any,
    tensor_var_names: Dict[str, str],
) -> None:
    if not model_file.exists():
        return
    desired_axis_by_first_output_var: Dict[str, int] = {}
    for op in list(model_ir.operators):
        if str(op.op_type) != "SPLIT" or len(op.outputs) == 0 or len(op.inputs) == 0:
            continue
        input_tensor = model_ir.tensors.get(str(op.inputs[-1]), None)
        if input_tensor is None:
            continue
        input_shape = [int(v) for v in list(input_tensor.shape)]
        if len(input_shape) == 0:
            continue
        output_shapes: list[list[int]] = []
        for output_name in list(op.outputs):
            output_tensor = model_ir.tensors.get(str(output_name), None)
            if output_tensor is None:
                output_shapes = []
                break
            output_shape = [int(v) for v in list(output_tensor.shape)]
            if len(output_shape) != len(input_shape):
                output_shapes = []
                break
            output_shapes.append(output_shape)
        if len(output_shapes) == 0:
            continue
        desired_axis = _resolve_static_split_axis(input_shape, output_shapes)
        if desired_axis is None:
            continue
        first_output_var = tensor_var_names.get(str(op.outputs[0]), "")
        if first_output_var != "":
            desired_axis_by_first_output_var[str(first_output_var)] = int(desired_axis)
    if len(desired_axis_by_first_output_var) == 0:
        return
    lines = model_file.read_text(encoding="utf-8").splitlines()
    split_pattern = re.compile(
        r"^(?P<indent>\s*)(?P<outputs>[A-Za-z0-9_, ]+)\s*=\s*list\(torch\.tensor_split\("
        r"(?P<input>[A-Za-z0-9_]+),\s*(?P<sections>\d+),\s*dim=_normalize_dim\((?P<axis>-?\d+),\s*"
        r"(?P=input)\.ndim\)\)\)\s*$"
    )
    rewritten = False
    for index, line in enumerate(lines):
        match = split_pattern.match(str(line))
        if match is None:
            continue
        output_vars = [token.strip() for token in str(match.group("outputs")).split(",") if token.strip() != ""]
        if len(output_vars) == 0:
            continue
        desired_axis = desired_axis_by_first_output_var.get(str(output_vars[0]), None)
        if desired_axis is None:
            continue
        current_axis = int(match.group("axis"))
        if int(current_axis) == int(desired_axis):
            continue
        lines[index] = (
            f"{match.group('indent')}{match.group('outputs')} = list(torch.tensor_split("
            f"{match.group('input')}, {match.group('sections')}, "
            f"dim=_normalize_dim({int(desired_axis)}, {match.group('input')}.ndim)))"
        )
        rewritten = True
    for index, line in enumerate(lines):
        match = split_pattern.match(str(line))
        if match is None:
            continue
        input_var = str(match.group("input"))
        current_axis = int(match.group("axis"))
        if not input_var.endswith("_cf") or int(current_axis) != 3:
            continue
        output_vars = [token.strip() for token in str(match.group("outputs")).split(",") if token.strip() != ""]
        if len(output_vars) == 0:
            continue
        lookahead = "\n".join(lines[index + 1:index + 9])
        concat_uses_outputs = any(f"{output_var}" in lookahead for output_var in output_vars)
        if (
            concat_uses_outputs
            and "_apply_concat(" in lookahead
            and "axis=3" in lookahead
        ):
            lines[index] = (
                f"{match.group('indent')}{match.group('outputs')} = list(torch.tensor_split("
                f"{match.group('input')}, {match.group('sections')}, "
                f"dim=_normalize_dim(1, {match.group('input')}.ndim)))"
            )
            rewritten = True
    if rewritten:
        model_file.write_text("\n".join(lines) + "\n", encoding="utf-8")


def _load_native_codegen_impl(bindings: _NativeCodegenBindings) -> Callable[..., Any]:
    global _NATIVE_CODEGEN_IMPL
    if bindings.compiled_impl is not None:
        return bindings.compiled_impl
    if _NATIVE_CODEGEN_IMPL is None:
        env: Dict[str, Any] = dict(bindings.module_globals)
        exec(_NATIVE_CODEGEN_FUNCTION_SOURCE, env, env)
        _NATIVE_CODEGEN_IMPL = env['_write_native_model_file_codegen_core_body_main_inner_legacy_impl']
    assert _NATIVE_CODEGEN_IMPL is not None
    bindings.compiled_impl = _NATIVE_CODEGEN_IMPL
    return cast(Callable[..., Any], _NATIVE_CODEGEN_IMPL)


def initialize_codegen_phase(
    state: _NativeCodegenState,
    bindings: _NativeCodegenBindings,
) -> None:
    state.exec_env = dict(bindings.module_globals)


def collect_special_specs_phase(
    state: _NativeCodegenState,
    bindings: _NativeCodegenBindings,
) -> None:
    _ = state
    bindings.compiled_impl = _load_native_codegen_impl(bindings)


def register_direct_modules_phase(
    state: _NativeCodegenState,
    bindings: _NativeCodegenBindings,
) -> None:
    _ = state
    _ = bindings


def emit_operator_forward_phase(
    state: _NativeCodegenState,
    bindings: _NativeCodegenBindings,
) -> None:
    _ = state
    _ = bindings


def postprocess_forward_phase(
    state: _NativeCodegenState,
    bindings: _NativeCodegenBindings,
) -> None:
    _ = state
    _ = bindings


def assemble_and_write_model_phase(
    state: _NativeCodegenState,
    bindings: _NativeCodegenBindings,
) -> None:
    impl = _load_native_codegen_impl(bindings)
    load_specs_result = impl(state.context)
    _rewrite_public_layout_bridge_binary_align_calls(
        state.context.package_dir / "model.py"
    )
    _rewrite_split_axes_from_static_shapes(
        state.context.package_dir / "model.py",
        model_ir=state.context.model_ir,
        tensor_var_names=state.context.tensor_var_names,
    )
    state.load_specs_result = [
        (str(attr_path), str(tensor_name))
        for attr_path, tensor_name in list(load_specs_result)
    ]


def execute_native_codegen_pipeline(
    state: _NativeCodegenState,
    bindings: _NativeCodegenBindings,
) -> None:
    initialize_codegen_phase(state, bindings)
    collect_special_specs_phase(state, bindings)
    register_direct_modules_phase(state, bindings)
    emit_operator_forward_phase(state, bindings)
    postprocess_forward_phase(state, bindings)
    assemble_and_write_model_phase(state, bindings)
