onnx2tf 1.29.18__py3-none-any.whl → 1.29.20__py3-none-any.whl

onnx2tf/ops/DeformConv.py

@@ -0,0 +1,399 @@
+import sys
+import random
+random.seed(0)
+import numpy as np
+np.random.seed(0)
+import tensorflow as tf
+import onnx_graphsurgeon as gs
+from onnx2tf.utils.common_functions import (
+    get_constant_or_variable,
+    get_weights_constant_or_variable,
+    print_node_info,
+    inverted_operation_enable_disable,
+    make_tf_node_info,
+    get_replacement_parameter,
+    pre_process_transpose,
+    post_process_transpose,
+    transpose_with_flexing_deterrence,
+)
+from onnx2tf.utils.logging import *
+
+INF_INDEX_VALUE: int = 4294967296
+
+
+def _to_int_tensor(value, name=None):
+    if isinstance(value, tf.Tensor):
+        return tf.cast(value, tf.int32)
+    return tf.constant(value, dtype=tf.int32, name=name)
+
+
+def _bilinear_sample_2d(
+    image,
+    coords,
+):
+    """
+    image: [N, H, W, C]
+    coords: [N, oH, oW, kH, kW, 2] in absolute coords (y, x)
+    """
+    coord_dtype = coords.dtype
+    h = tf.shape(image)[1]
+    w = tf.shape(image)[2]
+    h_f = tf.cast(h, coord_dtype)
+    w_f = tf.cast(w, coord_dtype)
+    max_y = h_f - 1.0
+    max_x = w_f - 1.0
+
+    y, x = tf.split(coords, num_or_size_splits=2, axis=-1)
+
+    y0 = tf.floor(y)
+    x0 = tf.floor(x)
+    y1 = y0 + 1.0
+    x1 = x0 + 1.0
+
+    dy = y - y0
+    dx = x - x0
+
+    w00 = (1.0 - dy) * (1.0 - dx)
+    w10 = dy * (1.0 - dx)
+    w11 = dy * dx
+    w01 = (1.0 - dy) * dx
+
+    def _in_bounds(y_idx, x_idx):
+        return tf.logical_and(
+            tf.logical_and(y_idx >= 0.0, y_idx <= max_y),
+            tf.logical_and(x_idx >= 0.0, x_idx <= max_x),
+        )
+
+    m00 = _in_bounds(y0, x0)
+    m10 = _in_bounds(y1, x0)
+    m11 = _in_bounds(y1, x1)
+    m01 = _in_bounds(y0, x1)
+
+    y0c = tf.clip_by_value(y0, 0.0, max_y)
+    x0c = tf.clip_by_value(x0, 0.0, max_x)
+    y1c = tf.clip_by_value(y1, 0.0, max_y)
+    x1c = tf.clip_by_value(x1, 0.0, max_x)
+
+    y0i = tf.cast(y0c, tf.int32)
+    x0i = tf.cast(x0c, tf.int32)
+    y1i = tf.cast(y1c, tf.int32)
+    x1i = tf.cast(x1c, tf.int32)
+
+    input_flat = tf.reshape(image, tf.stack([tf.shape(image)[0], h * w, tf.shape(image)[3]]))
+
+    def _gather(y_idx, x_idx):
+        linear = y_idx * w + x_idx
+        linear = tf.squeeze(linear, axis=-1)
+        return tf.gather(input_flat, linear, batch_dims=1)
+
+    v00 = _gather(y0i, x0i)
+    v10 = _gather(y1i, x0i)
+    v11 = _gather(y1i, x1i)
+    v01 = _gather(y0i, x1i)
+
+    m00 = tf.cast(m00, image.dtype)
+    m10 = tf.cast(m10, image.dtype)
+    m11 = tf.cast(m11, image.dtype)
+    m01 = tf.cast(m01, image.dtype)
+
+    output = w00 * m00 * v00 + w10 * m10 * v10 + w11 * m11 * v11 + w01 * m01 * v01
+    return output
+
+
+@print_node_info
+@inverted_operation_enable_disable
+@get_replacement_parameter
+def make_node(
+    *,
+    graph_node: gs.Node,
+    tf_layers_dict: dict,
+    **kwargs: dict,
+):
+    """DeformConv
+
+    Parameters
+    ----------
+    graph_node: gs.Node
+        graph_surgeon Node
+
+    tf_layers_dict: dict
+        optype, shape, dtype, tensorflow graph
+    """
+    before_op_output_shape_trans_1 = \
+        tf_layers_dict.get(graph_node.inputs[0].name, {}).get('before_op_output_shape_trans', True)
+    before_op_output_shape_trans_3 = \
+        tf_layers_dict.get(graph_node.inputs[2].name, {}).get('before_op_output_shape_trans', True)
+    before_op_output_shape_trans_4 = \
+        tf_layers_dict.get(graph_node.inputs[3].name, {}).get('before_op_output_shape_trans', True) \
+            if len(graph_node.inputs) >= 4 else True
+    before_op_output_shape_trans_5 = \
+        tf_layers_dict.get(graph_node.inputs[4].name, {}).get('before_op_output_shape_trans', True) \
+            if len(graph_node.inputs) >= 5 else True
+
+    graph_node_input_1 = get_constant_or_variable(
+        graph_node.inputs[0],
+        before_op_output_shape_trans_1,
+    )
+
+    kernel_shape = graph_node.attrs.get('kernel_shape', [])
+    if kernel_shape == [] and graph_node.inputs[1].shape is not None:
+        kernel_shape = graph_node.inputs[1].shape[2:]
+    kernel_size = len(kernel_shape) if kernel_shape != [] else 2
+
+    graph_node_input_2 = get_weights_constant_or_variable(
+        const_or_var=graph_node.inputs[1],
+        kernel_size=kernel_size,
+    )
+    graph_node_input_3 = get_constant_or_variable(
+        graph_node.inputs[2],
+        before_op_output_shape_trans_3,
+    )
+    graph_node_input_4 = get_constant_or_variable(
+        graph_node.inputs[3],
+        before_op_output_shape_trans_4,
+    ) if len(graph_node.inputs) >= 4 else None
+    graph_node_input_5 = get_constant_or_variable(
+        graph_node.inputs[4],
+        before_op_output_shape_trans_5,
+    ) if len(graph_node.inputs) >= 5 else None
+
+    graph_node_output: gs.Variable = graph_node.outputs[0]
+    output_tensor_shape = graph_node_output.shape
+    dtype = graph_node_output.dtype
+
+    input_tensor = tf_layers_dict[graph_node_input_1.name]['tf_node'] \
+        if isinstance(graph_node_input_1, gs.Variable) else graph_node_input_1
+    weights = tf_layers_dict[graph_node_input_2.name]['tf_node'] \
+        if isinstance(graph_node_input_2, gs.Variable) else graph_node_input_2
+    offset = tf_layers_dict[graph_node_input_3.name]['tf_node'] \
+        if isinstance(graph_node_input_3, gs.Variable) else graph_node_input_3
+    bias = tf_layers_dict[graph_node_input_4.name]['tf_node'] \
+        if isinstance(graph_node_input_4, gs.Variable) else graph_node_input_4
+    mask = tf_layers_dict[graph_node_input_5.name]['tf_node'] \
+        if isinstance(graph_node_input_5, gs.Variable) else graph_node_input_5
+
+    input_tensor_shape = input_tensor.shape
+
+    if input_tensor_shape is not None and len(input_tensor_shape) != 4:
+        error('DeformConv currently supports only 2D inputs (N, C, H, W).')
+        sys.exit(1)
+
+    # Preserving Graph Structure (Dict)
+    tf_layers_dict[graph_node_output.name] = {
+        'optype': graph_node.op,
+        'shape': output_tensor_shape,
+        'dtype': dtype,
+        'nhwc': True,
+    }
+
+    # Pre-process transpose
+    input_tensor = pre_process_transpose(
+        value_before_transpose=input_tensor,
+        param_target='inputs',
+        param_name=graph_node.inputs[0].name,
+        **kwargs,
+    )
+    offset = pre_process_transpose(
+        value_before_transpose=offset,
+        param_target='inputs',
+        param_name=graph_node.inputs[2].name,
+        **kwargs,
+    )
+    if mask is not None:
+        mask = pre_process_transpose(
+            value_before_transpose=mask,
+            param_target='inputs',
+            param_name=graph_node.inputs[4].name,
+            **kwargs,
+        )
+
+    input_dtype = input_tensor.dtype
+    if weights is not None and weights.dtype != input_dtype:
+        weights = tf.cast(weights, input_dtype)
+    if offset is not None and offset.dtype != input_dtype:
+        offset = tf.cast(offset, input_dtype)
+    if bias is not None and bias.dtype != input_dtype:
+        bias = tf.cast(bias, input_dtype)
+    if mask is not None and mask.dtype != input_dtype:
+        mask = tf.cast(mask, input_dtype)
+
+    # Workaround to avoid as many conversion failures as possible
+    onnx_input_shape = [
+        dim if isinstance(dim, int) else None for dim in graph_node.inputs[0].shape
+    ] if graph_node.inputs[0].shape is not None else None
+    tf_input_shape = [
+        dim if isinstance(dim, int) else None for dim in input_tensor.shape
+    ]
+    if onnx_input_shape is not None \
+        and len(onnx_input_shape) > 1 and len(tf_input_shape) > 1 \
+        and onnx_input_shape == tf_input_shape:
+
+        shape_for_judging_skip = [
+            dim if dim is not None else INF_INDEX_VALUE for dim in onnx_input_shape[1:]
+        ]
+        if shape_for_judging_skip.count(shape_for_judging_skip[0]) != len(shape_for_judging_skip):
+            input_tensor = transpose_with_flexing_deterrence(
+                input_tensor=input_tensor,
+                perm=[0,2,3,1],
+                **kwargs,
+            )
+            offset = transpose_with_flexing_deterrence(
+                input_tensor=offset,
+                perm=[0,2,3,1],
+                **kwargs,
+            )
+            if mask is not None:
+                mask = transpose_with_flexing_deterrence(
+                    input_tensor=mask,
+                    perm=[0,2,3,1],
+                    **kwargs,
+                )
+
+    # Attributes
+    dilations = graph_node.attrs.get('dilations', [1, 1])
+    group = graph_node.attrs.get('group', 1)
+    offset_group = graph_node.attrs.get('offset_group', 1)
+    pads = graph_node.attrs.get('pads', [0, 0, 0, 0])
+    strides = graph_node.attrs.get('strides', [1, 1])
+
+    dilation_h, dilation_w = dilations
+    stride_h, stride_w = strides
+    pad_top, pad_left, pad_bottom, pad_right = pads
+
+    # Input prep
+    if pad_top != 0 or pad_bottom != 0 or pad_left != 0 or pad_right != 0:
+        input_tensor = tf.pad(
+            input_tensor,
+            paddings=[[0, 0], [pad_top, pad_bottom], [pad_left, pad_right], [0, 0]],
+        )
+
+    batch = tf.shape(input_tensor)[0]
+    in_h = tf.shape(input_tensor)[1]
+    in_w = tf.shape(input_tensor)[2]
+    in_c = tf.shape(input_tensor)[3]
+
+    offset_shape = tf.shape(offset)
+    out_h = offset_shape[1]
+    out_w = offset_shape[2]
+
+    # Kernel shape
+    if kernel_shape != []:
+        kh = _to_int_tensor(kernel_shape[0])
+        kw = _to_int_tensor(kernel_shape[1])
+    else:
+        kh = _to_int_tensor(tf.shape(weights)[0])
+        kw = _to_int_tensor(tf.shape(weights)[1])
+
+    # Base grid: [oH, oW, kH, kW, 2]
+    oy = tf.range(out_h, dtype=input_dtype) * tf.cast(stride_h, input_dtype)
+    ox = tf.range(out_w, dtype=input_dtype) * tf.cast(stride_w, input_dtype)
+    ky = tf.range(kh, dtype=input_dtype) * tf.cast(dilation_h, input_dtype)
+    kx = tf.range(kw, dtype=input_dtype) * tf.cast(dilation_w, input_dtype)
+
+    oy = tf.reshape(oy, tf.stack([out_h, 1, 1, 1]))
+    ox = tf.reshape(ox, tf.stack([1, out_w, 1, 1]))
+    ky = tf.reshape(ky, tf.stack([1, 1, kh, 1]))
+    kx = tf.reshape(kx, tf.stack([1, 1, 1, kw]))
+
+    y = oy + ky
+    x = ox + kx
+    target_shape = tf.stack([out_h, out_w, kh, kw])
+    y = tf.broadcast_to(y, target_shape)
+    x = tf.broadcast_to(x, target_shape)
+    base_grid = tf.stack([y, x], axis=-1)
+
+    # Offset reshape: [N, oH, oW, Goff, kH, kW, 2]
+    offset = tf.reshape(
+        offset,
+        tf.stack([batch, out_h, out_w, offset_group, kh, kw, 2]),
+    )
+
+    coords = base_grid[None, :, :, None, :, :, :] + offset
+    coords = tf.transpose(coords, [0, 3, 1, 2, 4, 5, 6])
+    coords = tf.reshape(coords, tf.stack([batch * offset_group, out_h, out_w, kh, kw, 2]))
+
+    # Input grouping for offset_group
+    c_per_offset = tf.math.floordiv(in_c, offset_group)
+    input_tensor = tf.reshape(
+        input_tensor,
+        tf.stack([batch, in_h, in_w, offset_group, c_per_offset]),
+    )
+    input_tensor = tf.transpose(input_tensor, [0, 3, 1, 2, 4])
+    input_tensor = tf.reshape(
+        input_tensor,
+        tf.stack([batch * offset_group, in_h, in_w, c_per_offset]),
+    )
+
+    sampled = _bilinear_sample_2d(input_tensor, coords)
+    sampled = tf.reshape(
+        sampled,
+        tf.stack([batch, offset_group, out_h, out_w, kh, kw, c_per_offset]),
+    )
+    sampled = tf.transpose(sampled, [0, 2, 3, 1, 4, 5, 6])
+
+    if mask is not None:
+        mask = tf.reshape(
+            mask,
+            tf.stack([batch, out_h, out_w, offset_group, kh, kw, 1]),
+        )
+        sampled = sampled * tf.cast(mask, sampled.dtype)
+
+    # Merge offset_group back to channel dim: [N, oH, oW, kH, kW, C]
+    sampled = tf.reshape(
+        sampled,
+        tf.stack([batch, out_h, out_w, kh, kw, in_c]),
+    )
+
+    # Grouped convolution via batched matmul
+    out_c = tf.shape(weights)[3]
+    c_per_group = tf.math.floordiv(in_c, group)
+    out_c_per_group = tf.math.floordiv(out_c, group)
+
+    cols = tf.reshape(sampled, tf.stack([batch * out_h * out_w, kh * kw * in_c]))
+    cols = tf.reshape(cols, tf.stack([batch * out_h * out_w, group, kh * kw * c_per_group]))
+    cols = tf.transpose(cols, [1, 0, 2])
+
+    weights = tf.reshape(weights, tf.stack([kh, kw, c_per_group, group, out_c_per_group]))
+    weights = tf.transpose(weights, [3, 0, 1, 2, 4])
+    weights = tf.reshape(weights, tf.stack([group, kh * kw * c_per_group, out_c_per_group]))
+
+    output = tf.matmul(cols, weights)
+    output = tf.transpose(output, [1, 0, 2])
+    output = tf.reshape(output, tf.stack([batch, out_h, out_w, out_c]))
+
+    if bias is not None:
+        output += tf.reshape(bias, tf.stack([1, 1, 1, out_c]))
+
+    if output.dtype != input_dtype:
+        output = tf.cast(output, input_dtype)
+
+    # Post-process transpose
+    tf_layers_dict[graph_node_output.name]['tf_node'] = post_process_transpose(
+        value_before_transpose=output,
+        param_target='outputs',
+        param_name=graph_node.outputs[0].name,
+        **kwargs,
+    )
+
+    # Generation of Debug Info
+    tf_layers_dict[graph_node_output.name]['tf_node_info'] = \
+        make_tf_node_info(
+            node_info={
+                'tf_op_type': 'DeformConv',
+                'tf_inputs': {
+                    'input_tensor': input_tensor,
+                    'weights': weights,
+                    'offset': offset,
+                    'bias': bias,
+                    'mask': mask,
+                    'strides': strides,
+                    'dilations': dilations,
+                    'pads': pads,
+                    'group': group,
+                    'offset_group': offset_group,
+                },
+                'tf_outputs': {
+                    'output': tf_layers_dict[graph_node_output.name]['tf_node'],
+                },
+            }
+        )

onnx2tf/ops/GatherElements.py

@@ -57,9 +57,10 @@ def make_node(
         graph_node.inputs[0],
         before_op_output_shape_trans,
     )
+    # Indices must not be layout-transposed.
     graph_node_input_2 = get_constant_or_variable(
         graph_node.inputs[1],
-        before_op_output_shape_trans,
+        False,
     )
     graph_node_output: gs.Variable = graph_node.outputs[0]
     shape = graph_node_output.shape
@@ -77,12 +78,29 @@ def make_node(
         param_name=graph_node.inputs[0].name,
         **kwargs,
     )
-    indices_tensor = pre_process_transpose(
-        value_before_transpose=indices_tensor,
-        param_target='inputs',
-        param_name=graph_node.inputs[1].name,
-        **kwargs,
-    )
+    # If input is transposed by replacement params, align indices tensor shape.
+    op_rep_params = kwargs.get('op_rep_params', [])
+    params_perm = None
+    indices_perm = None
+    for op_rep_param in op_rep_params:
+        if op_rep_param['param_target'] == 'inputs' \
+            and op_rep_param['param_name'] == graph_node.inputs[0].name:
+            params_perm = op_rep_param.get('pre_process_transpose_perm', None)
+        if op_rep_param['param_target'] == 'inputs' \
+            and op_rep_param['param_name'] == graph_node.inputs[1].name:
+            indices_perm = op_rep_param.get('pre_process_transpose_perm', None)
+    target_perm = indices_perm if indices_perm is not None else params_perm
+    if target_perm is not None:
+        try:
+            rank = len(indices_tensor.shape) if hasattr(indices_tensor, "shape") else None
+            if rank is None or rank == len(target_perm):
+                indices_tensor = transpose_with_flexing_deterrence(
+                    input_tensor=indices_tensor,
+                    perm=target_perm,
+                    **kwargs,
+                )
+        except Exception:
+            pass
 
     tensor_rank = len(input_tensor.shape)
 

onnx2tf/ops/GatherND.py

@@ -51,9 +51,10 @@ def make_node(
         graph_node.inputs[0],
         before_op_output_shape_trans,
     )
+    # Indices must not be layout-transposed.
     graph_node_input_2 = get_constant_or_variable(
         graph_node.inputs[1],
-        before_op_output_shape_trans,
+        False,
     )
     graph_node_output: gs.Variable = graph_node.outputs[0]
     shape = graph_node_output.shape
@@ -89,6 +90,32 @@ def make_node(
 
     replace_gathernd_to_pseudo_gathernd = "gathernd" in kwargs['replace_to_pseudo_operators']
 
+    # If params is transposed, adjust indices to match the transposed layout.
+    op_rep_params = kwargs.get('op_rep_params', [])
+    params_perm = None
+    indices_perm_specified = False
+    for op_rep_param in op_rep_params:
+        if op_rep_param['param_target'] == 'inputs' and op_rep_param['param_name'] == graph_node.inputs[0].name:
+            params_perm = op_rep_param.get('pre_process_transpose_perm', None)
+        if op_rep_param['param_target'] == 'inputs' and op_rep_param['param_name'] == graph_node.inputs[1].name:
+            if op_rep_param.get('pre_process_transpose_perm', None) is not None:
+                indices_perm_specified = True
+    if params_perm is not None and not indices_perm_specified:
+        # Only handle standard layout swaps that keep batch dims at the front.
+        if batch_dims <= len(params_perm) \
+            and list(params_perm[:batch_dims]) == list(range(batch_dims)):
+            perm_tail = [p - batch_dims for p in params_perm if p >= batch_dims]
+            try:
+                if isinstance(indices_tensor, np.ndarray):
+                    if indices_tensor.shape and indices_tensor.shape[-1] == len(perm_tail):
+                        indices_tensor = indices_tensor[..., perm_tail]
+                else:
+                    idx_last = indices_tensor.shape[-1] if indices_tensor.shape is not None else None
+                    if idx_last is None or idx_last == len(perm_tail):
+                        indices_tensor = tf.gather(indices_tensor, perm_tail, axis=-1)
+            except Exception:
+                pass
+
     # Preserving Graph Structure (Dict)
     tf_layers_dict[graph_node_output.name] = {
         'optype': graph_node.op,

onnx2tf/ops/ScatterElements.py

@@ -55,9 +55,10 @@ def make_node(
         graph_node.inputs[0],
         before_op_output_shape_trans,
     )
+    # Indices must not be layout-transposed.
     graph_node_input_2 = get_constant_or_variable(
         graph_node.inputs[1],
-        before_op_output_shape_trans,
+        False,
     )
     graph_node_input_3 = get_constant_or_variable(
         graph_node.inputs[2],
@@ -81,12 +82,29 @@ def make_node(
     indices_tensor = tf_layers_dict[graph_node_input_2.name]['tf_node'] \
         if isinstance(graph_node_input_2, gs.Variable) else graph_node_input_2
     # Pre-process transpose
-    indices_tensor = pre_process_transpose(
-        value_before_transpose=indices_tensor,
-        param_target='inputs',
-        param_name=graph_node.inputs[1].name,
-        **kwargs,
-    )
+    # If input is transposed by replacement params, align indices tensor shape.
+    op_rep_params = kwargs.get('op_rep_params', [])
+    params_perm = None
+    indices_perm = None
+    for op_rep_param in op_rep_params:
+        if op_rep_param['param_target'] == 'inputs' \
+            and op_rep_param['param_name'] == graph_node.inputs[0].name:
+            params_perm = op_rep_param.get('pre_process_transpose_perm', None)
+        if op_rep_param['param_target'] == 'inputs' \
+            and op_rep_param['param_name'] == graph_node.inputs[1].name:
+            indices_perm = op_rep_param.get('pre_process_transpose_perm', None)
+    target_perm = indices_perm if indices_perm is not None else params_perm
+    if target_perm is not None:
+        try:
+            rank = len(indices_tensor.shape) if hasattr(indices_tensor, "shape") else None
+            if rank is None or rank == len(target_perm):
+                indices_tensor = transpose_with_flexing_deterrence(
+                    input_tensor=indices_tensor,
+                    perm=target_perm,
+                    **kwargs,
+                )
+        except Exception:
+            pass
     updates_tensor = tf_layers_dict[graph_node_input_3.name]['tf_node'] \
         if isinstance(graph_node_input_3, gs.Variable) else graph_node_input_3
     # Pre-process transpose

onnx2tf/ops/ScatterND.py

@@ -13,6 +13,7 @@ from onnx2tf.utils.common_functions import (
     get_replacement_parameter,
     pre_process_transpose,
     post_process_transpose,
+    transpose_with_flexing_deterrence,
 )
 
 
@@ -79,6 +80,32 @@ def make_node(
                 and 'nhwc' in tf_layers_dict[graph_node_input_1.name].keys() else False
     }
 
+    op_rep_params = kwargs.get('op_rep_params', [])
+    params_perm = None
+    indices_perm = None
+    for op_rep_param in op_rep_params:
+        if op_rep_param['param_target'] == 'inputs' \
+            and op_rep_param['param_name'] == graph_node.inputs[0].name:
+            params_perm = op_rep_param.get('pre_process_transpose_perm', None)
+        if op_rep_param['param_target'] == 'inputs' \
+            and op_rep_param['param_name'] == graph_node.inputs[1].name:
+            indices_perm = op_rep_param.get('pre_process_transpose_perm', None)
+
+    def reorder_indices_last_dim(target_indices, perm):
+        if perm is None:
+            return target_indices
+        try:
+            if isinstance(target_indices, np.ndarray):
+                if target_indices.shape and target_indices.shape[-1] == len(perm):
+                    return target_indices[..., perm]
+            else:
+                idx_last = target_indices.shape[-1] if target_indices.shape is not None else None
+                if idx_last is None or idx_last == len(perm):
+                    return tf.gather(target_indices, perm, axis=-1)
+        except Exception:
+            pass
+        return target_indices
+
     # Pre-process transpose
     input_tensor = pre_process_transpose(
         value_before_transpose=input_tensor,
@@ -86,18 +113,26 @@ def make_node(
         param_name=graph_node.inputs[0].name,
         **kwargs,
     )
-    indices_tensor = pre_process_transpose(
-        value_before_transpose=indices_tensor,
-        param_target='inputs',
-        param_name=graph_node.inputs[1].name,
-        **kwargs,
-    )
+    # Indices must not be layout-transposed; apply explicit perm only if specified.
+    if indices_perm is not None:
+        try:
+            rank = len(indices_tensor.shape) if hasattr(indices_tensor, "shape") else None
+            if rank is None or rank == len(indices_perm):
+                indices_tensor = transpose_with_flexing_deterrence(
+                    input_tensor=indices_tensor,
+                    perm=indices_perm,
+                    **kwargs,
+                )
+        except Exception:
+            pass
     updates_tensor = pre_process_transpose(
         value_before_transpose=updates_tensor,
         param_target='inputs',
         param_name=graph_node.inputs[2].name,
         **kwargs,
     )
+    if params_perm is not None and indices_perm is None:
+        indices_tensor = reorder_indices_last_dim(indices_tensor, params_perm)
 
     # When NHWC is fixed, return to NCHW format before processing.
     data_nhwc = tf_layers_dict[graph_node_input_1.name]['nhwc'] \
@@ -119,6 +154,8 @@ def make_node(
         and len(input_tensor.shape) >= 3:
         perm = [0, len(input_tensor.shape)-1] + [i for i in range(1, len(input_tensor.shape)-1)]
         input_tensor = tf.transpose(a=input_tensor, perm=perm)
+        if indices_perm is None:
+            indices_tensor = reorder_indices_last_dim(indices_tensor, perm)
         nchw = True
     elif not data_nhwc \
         and len(input_tensor.shape) >= 3 \
@@ -126,6 +163,8 @@ def make_node(
         and input_tensor.shape != graph_node.inputs[0].shape:
         perm = [0, len(input_tensor.shape)-1] + [i for i in range(1, len(input_tensor.shape)-1)]
         input_tensor = tf.transpose(a=input_tensor, perm=perm)
+        if indices_perm is None:
+            indices_tensor = reorder_indices_last_dim(indices_tensor, perm)
         nchw = True
     ## indices
     if indices_nhwc \

onnx2tf/ops/TensorScatter.py

@@ -14,6 +14,7 @@ from onnx2tf.utils.common_functions import (
     get_replacement_parameter,
     pre_process_transpose,
     post_process_transpose,
+    transpose_with_flexing_deterrence,
 )
 from onnx2tf.utils.enums import NUMPY_DTYPES_TO_TF_DTYPES
 from onnx2tf.utils.logging import *
@@ -112,12 +113,25 @@ def make_node(
         **kwargs,
     )
     if write_indices is not None:
-        write_indices = pre_process_transpose(
-            value_before_transpose=write_indices,
-            param_target='inputs',
-            param_name=graph_node.inputs[2].name,
-            **kwargs,
-        )
+        # Indices must not be layout-transposed; apply explicit perm only if specified.
+        op_rep_params = kwargs.get('op_rep_params', [])
+        indices_perm = None
+        for op_rep_param in op_rep_params:
+            if op_rep_param['param_target'] == 'inputs' \
+                and op_rep_param['param_name'] == graph_node.inputs[2].name:
+                indices_perm = op_rep_param.get('pre_process_transpose_perm', None)
+                break
+        if indices_perm is not None:
+            try:
+                rank = len(write_indices.shape) if hasattr(write_indices, "shape") else None
+                if rank is None or rank == len(indices_perm):
+                    write_indices = transpose_with_flexing_deterrence(
+                        input_tensor=write_indices,
+                        perm=indices_perm,
+                        **kwargs,
+                    )
+            except Exception:
+                pass
 
     # Generation of TF OP
     past_cache = _as_tensor(past_cache)