onnx2tf 1.29.15__py3-none-any.whl → 1.29.17__py3-none-any.whl

onnx2tf/ops/QLinearConv.py

@@ -1,4 +1,3 @@
-import sys
 import random
 random.seed(0)
 import numpy as np
@@ -25,45 +24,57 @@ def _dequantize_tensor(
     # Do computation in float32
     base = tf.cast(base, tf.float32)
     zero_point = tf.cast(zero_point, tf.float32)
+    scale = tf.cast(scale, tf.float32)
     return (base - zero_point) * scale
 
 
+def _reshape_per_output_channel(
+    *,
+    value,
+    weights,
+):
+    value_rank = len(value.shape)
+    weights_rank = len(weights.shape)
+    if value_rank == 0:
+        return value
+    if value_rank == 1 and weights_rank is not None:
+        shape = [1] * weights_rank
+        shape[-1] = -1
+        return tf.reshape(value, shape)
+    return value
+
+
 def _dequantize_weights(
     *,
     base,
     zero_point,
     scale,
-    is_bias=False,
-    scale_is_scalar=False,
 ):
     # Do computation in float32
     casted_base = tf.cast(base, tf.float32)
     casted_zero_point = tf.cast(zero_point, tf.float32)
-    spartial_shape_len = len(casted_base.shape) - 2
-    casted_zero_point_shape = casted_zero_point.shape[0]
-    if casted_zero_point_shape == base.shape[-2]:
-        reshaped_zero_point = tf.reshape(
-            tensor=casted_zero_point,
-            shape=[1 for _ in range(spartial_shape_len)] + [casted_zero_point_shape, 1],
-        )
-        if scale_is_scalar:
-            reshaped_scale = tf.reshape(
-                tensor=scale,
-                shape=[1 for _ in range(spartial_shape_len)] + [casted_zero_point_shape, 1],
-            )
-            tensor_list = [
-                (casted_base[..., i:i+1] - reshaped_zero_point) * reshaped_scale
-                for i in range(base.shape[-1])
-            ]
-            out_tensor = tf.concat(tensor_list, axis=-1)
-        else:
-            reshaped_scale = scale
-            out_tensor = (casted_base - reshaped_zero_point) * reshaped_scale
-        return tf.reshape(out_tensor, base.shape)
-    else:
-        reshaped_zero_point = casted_zero_point
-        reshaped_scale = scale
-        return (casted_base - reshaped_zero_point) * reshaped_scale
+    casted_scale = tf.cast(scale, tf.float32)
+    casted_zero_point = _reshape_per_output_channel(
+        value=casted_zero_point,
+        weights=casted_base,
+    )
+    casted_scale = _reshape_per_output_channel(
+        value=casted_scale,
+        weights=casted_base,
+    )
+    return (casted_base - casted_zero_point) * casted_scale
+
+
+def _get_qmin_qmax(dtype: tf.dtypes.DType):
+    if dtype == tf.uint8:
+        return 0.0, 255.0
+    if dtype == tf.int8:
+        return -128.0, 127.0
+    if dtype == tf.uint16:
+        return 0.0, 65535.0
+    if dtype == tf.int16:
+        return -32768.0, 32767.0
+    return None, None
 
 
 @print_node_info
@@ -139,6 +150,11 @@ def make_node(
 
     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
+    input_is_dequantized = False
+    input_nhwc = False
+    if isinstance(graph_node_input_1, gs.Variable):
+        input_is_dequantized = tf_layers_dict.get(graph_node_input_1.name, {}).get('is_dequantized', False)
+        input_nhwc = tf_layers_dict.get(graph_node_input_1.name, {}).get('nhwc', False)
     input_tensor_scale = tf_layers_dict[graph_node_input_2.name]['tf_node'] \
         if isinstance(graph_node_input_2, gs.Variable) else graph_node_input_2
     input_tensor_zero_point = tf_layers_dict[graph_node_input_3.name]['tf_node'] \
@@ -155,7 +171,7 @@ def make_node(
         if isinstance(graph_node_input_8, gs.Variable) else graph_node_input_8
     input_bias = tf_layers_dict[graph_node_input_9.name]['tf_node'] \
         if isinstance(graph_node_input_9, gs.Variable) else graph_node_input_9
-    output_dtype = input_tensor.dtype if input_tensor.dtype not in [tf.int8, tf.uint8] else tf.float32
+    output_quant_dtype = y_zero_point.dtype
 
     input_tensor_shape = input_tensor.shape
     input_tensor_rank = len(input_tensor_shape)
@@ -172,48 +188,32 @@ def make_node(
         'optype': graph_node.op,
         'shape': output_tensor_shape,
         'dtype': dtype,
+        'is_dequantized': True,
+        'nhwc': input_nhwc,
     }
 
     # Generation of TF OP
 
-    # Convert w_zero_point and w_scale to 1-D if scalar
-    if len(input_weights_zero_point.shape) == 0:
-        input_weights_zero_point = tf.fill([input_tensor.shape[-1]//group], input_weights_zero_point)
-    elif len(input_weights_zero_point.shape) > 1:
-        error(
-            f'Unsupported zero point: {graph_node.name} {input_weights_zero_point}'
-        )
-        sys.exit(1)
-
-    weights_scale_is_scalar = False
-    if len(input_weights_scale.shape) == 0:
-        weights_scale_is_scalar = True
-        input_weights_scale = tf.fill([input_tensor.shape[-1]//group], input_weights_scale)
-    elif len(input_weights_scale.shape) > 1:
-        error(
-            f'Unsupported scalet: {graph_node.name} {input_weights_scale}'
-        )
-        sys.exit(1)
-
     # Dequantize variables to float32
-    input_tensor = _dequantize_tensor(
-        base=input_tensor,
-        zero_point=input_tensor_zero_point,
-        scale=input_tensor_scale,
-    )
+    if input_is_dequantized:
+        input_tensor = tf.cast(input_tensor, tf.float32)
+    else:
+        input_tensor = _dequantize_tensor(
+            base=input_tensor,
+            zero_point=input_tensor_zero_point,
+            scale=input_tensor_scale,
+        )
     input_weights = _dequantize_weights(
         base=input_weights,
         zero_point=input_weights_zero_point,
         scale=input_weights_scale,
-        scale_is_scalar=weights_scale_is_scalar,
     )
-    y_zero_point = tf.cast(y_zero_point, tf.float32)
 
     # if bias is defined save it here
     if input_bias is not None:
         input_bias = tf.cast(input_bias, tf.float32)
-        input_bias_scale = input_tensor_scale * input_weights_scale
-        input_bias = tf.round(input_bias / input_bias_scale)
+        input_bias_scale = tf.cast(input_tensor_scale, tf.float32) * tf.cast(input_weights_scale, tf.float32)
+        input_bias = input_bias * input_bias_scale
 
     """
     Conv1D
@@ -260,7 +260,7 @@ def make_node(
         depthwise = bool(group == input_tensor_shape[-1])
 
     if depthwise is True:
-        depthwise_filter_shape = list(input_weights_shape[0:2]) + [-1, input_weights_shape[3] // group]
+        depthwise_filter_shape = list(input_weights_shape[0:2]) + [input_weights_shape[2], input_weights_shape[3] // group]
         input_weights = tf.reshape(input_weights, depthwise_filter_shape)
 
     # Conv
@@ -308,27 +308,23 @@ def make_node(
             )
         tf_op_type = tf.nn.depthwise_conv2d
 
-    # Process output
-    scaled_conv_node = tf.add(
-        x=tf.round(
-            tf.divide(
-                x=conv_node,
-                y=y_scale,
-            ),
-        ),
-        y=y_zero_point,
-    )
-
-    # Add bias to the convolution
+    # Add bias to the convolution (float)
     if input_bias is not None:
-        scaled_conv_node = tf.add(
-            x=scaled_conv_node,
+        conv_node = tf.add(
+            x=conv_node,
             y=input_bias,
         )
 
-    casted_conv_node = tf.cast(scaled_conv_node, output_dtype)
+    # quantize then dequantize to float32
+    y_scale = tf.cast(y_scale, tf.float32)
+    y_zero_point = tf.cast(y_zero_point, tf.float32)
+    quantized = tf.round(tf.divide(conv_node, y_scale)) + y_zero_point
+    qmin, qmax = _get_qmin_qmax(output_quant_dtype)
+    if qmin is not None and qmax is not None:
+        quantized = tf.clip_by_value(quantized, qmin, qmax)
+    dequantized = tf.multiply(tf.subtract(quantized, y_zero_point), y_scale)
 
-    tf_layers_dict[graph_node_output.name]['tf_node'] = casted_conv_node
+    tf_layers_dict[graph_node_output.name]['tf_node'] = dequantized
 
     # Generation of Debug Info
     tf_layers_dict[graph_node_output.name]['tf_node_info'] = \
@@ -349,4 +345,3 @@ def make_node(
                 },
             }
         )
-

onnx2tf/ops/QLinearMatMul.py

@@ -11,6 +11,47 @@ from onnx2tf.utils.common_functions import (
     make_tf_node_info,
 )
 
+def _get_qmin_qmax(dtype: tf.dtypes.DType):
+    if dtype == tf.uint8:
+        return 0.0, 255.0
+    if dtype == tf.int8:
+        return -128.0, 127.0
+    if dtype == tf.uint16:
+        return 0.0, 65535.0
+    if dtype == tf.int16:
+        return -32768.0, 32767.0
+    return None, None
+
+
+def _reshape_for_axis(
+    *,
+    value,
+    input_tensor,
+    axis: int,
+):
+    value_rank = len(value.shape)
+    input_rank = len(input_tensor.shape)
+    if value_rank == 1 and input_rank is not None:
+        shape = [1] * input_rank
+        shape[axis] = -1
+        return tf.reshape(value, shape)
+    return value
+
+
+def _reshape_for_output(
+    *,
+    value,
+    output_tensor,
+):
+    value_rank = len(value.shape)
+    output_rank = len(output_tensor.shape)
+    if value_rank == 1 and output_rank is not None and output_rank >= 2:
+        if output_tensor.shape[-2] == value.shape[0]:
+            shape = [1] * output_rank
+            shape[-2] = -1
+            return tf.reshape(value, shape)
+    return value
+
 
 @print_node_info
 @inverted_operation_enable_disable
@@ -76,12 +117,18 @@ def make_node(
 
     a = tf_layers_dict[graph_node_input_1.name]['tf_node'] \
         if isinstance(graph_node_input_1, gs.Variable) else graph_node_input_1
+    a_is_dequantized = False
+    if isinstance(graph_node_input_1, gs.Variable):
+        a_is_dequantized = tf_layers_dict.get(graph_node_input_1.name, {}).get('is_dequantized', False)
     a_scale = tf_layers_dict[graph_node_input_2.name]['tf_node'] \
         if isinstance(graph_node_input_2, gs.Variable) else graph_node_input_2
     a_zero_point = tf_layers_dict[graph_node_input_3.name]['tf_node'] \
         if isinstance(graph_node_input_3, gs.Variable) else graph_node_input_3
     b = tf_layers_dict[graph_node_input_4.name]['tf_node'] \
         if isinstance(graph_node_input_4, gs.Variable) else graph_node_input_4
+    b_is_dequantized = False
+    if isinstance(graph_node_input_4, gs.Variable):
+        b_is_dequantized = tf_layers_dict.get(graph_node_input_4.name, {}).get('is_dequantized', False)
     b_scale = tf_layers_dict[graph_node_input_5.name]['tf_node'] \
         if isinstance(graph_node_input_5, gs.Variable) else graph_node_input_5
     b_zero_point = tf_layers_dict[graph_node_input_6.name]['tf_node'] \
@@ -90,50 +137,60 @@ def make_node(
         if isinstance(graph_node_input_7, gs.Variable) else graph_node_input_7
     y_zero_point = tf_layers_dict[graph_node_input_8.name]['tf_node'] \
         if isinstance(graph_node_input_8, gs.Variable) else graph_node_input_8
-    y_dtype = y_zero_point.dtype if y_zero_point.dtype not in [tf.int8, tf.uint8] else tf.float32
+    y_dtype = y_zero_point.dtype
 
     # Preserving Graph Structure (Dict)
     tf_layers_dict[graph_node_output.name] = {
         'optype': graph_node.op,
         'shape': shape,
         'dtype': dtype,
+        'is_dequantized': True,
     }
 
     # Generation of TF OP
 
-    # reshape 1-D a_scale, a_zero_point, y_scale and
-    # y_zero_point so it can broadcast in arithmetic
-    # operations later
-    a_scale_shape = a_scale.shape
-    if a_scale_shape and a_scale_shape[0] > 1:
-        a_scale = tf.reshape(a_scale, [a_scale_shape[0], 1])
-        a_zero_point = tf.reshape(a_zero_point, [a_scale_shape[0], 1])
-    y_scale_shape = y_scale.shape
-    if y_scale_shape and y_scale_shape[0] > 1:
-        y_scale = tf.reshape(y_scale, [y_scale_shape[0], 1])
-        y_zero_point = tf.reshape(y_zero_point, [y_scale_shape[0], 1])
+    # reshape a_scale and a_zero_point to broadcast on row axis (second last)
+    a_scale = _reshape_for_axis(value=a_scale, input_tensor=a, axis=-2)
+    a_zero_point = _reshape_for_axis(value=a_zero_point, input_tensor=a, axis=-2)
+    # reshape b_scale and b_zero_point to broadcast on column axis (last)
+    b_scale = _reshape_for_axis(value=b_scale, input_tensor=b, axis=-1)
+    b_zero_point = _reshape_for_axis(value=b_zero_point, input_tensor=b, axis=-1)
 
     # cast all inputs to float32
     a = tf.cast(a, tf.float32)
+    a_scale = tf.cast(a_scale, tf.float32)
     a_zero_point = tf.cast(a_zero_point, tf.float32)
     b = tf.cast(b, tf.float32)
+    b_scale = tf.cast(b_scale, tf.float32)
     b_zero_point = tf.cast(b_zero_point, tf.float32)
+    y_scale = tf.cast(y_scale, tf.float32)
     y_zero_point = tf.cast(y_zero_point, tf.float32)
 
     # dequantize a and b
-    dequantized_a = tf.subtract(a, a_zero_point)
-    dequantized_a = tf.multiply(dequantized_a, a_scale)
-    dequantized_b = tf.subtract(b, b_zero_point)
-    dequantized_b = tf.multiply(dequantized_b, b_scale)
+    if a_is_dequantized:
+        dequantized_a = tf.cast(a, tf.float32)
+    else:
+        dequantized_a = tf.multiply(tf.subtract(a, a_zero_point), a_scale)
+
+    if b_is_dequantized:
+        dequantized_b = tf.cast(b, tf.float32)
+    else:
+        dequantized_b = tf.multiply(tf.subtract(b, b_zero_point), b_scale)
 
     # matmul
     x = tf.matmul(dequantized_a, dequantized_b)
 
-    # quantize x
-    y = tf.divide(x, y_scale)
-    y = tf.round(y)
+    # broadcast output scale/zero_point if needed
+    y_scale = _reshape_for_output(value=y_scale, output_tensor=x)
+    y_zero_point = _reshape_for_output(value=y_zero_point, output_tensor=x)
+
+    # quantize then dequantize to float32
+    y = tf.round(tf.divide(x, y_scale))
     y = tf.add(y, y_zero_point)
-    y = tf.saturate_cast(y, y_dtype)
+    qmin, qmax = _get_qmin_qmax(y_dtype)
+    if qmin is not None and qmax is not None:
+        y = tf.clip_by_value(y, qmin, qmax)
+    y = tf.multiply(tf.subtract(y, y_zero_point), y_scale)
 
     tf_layers_dict[graph_node_output.name]['tf_node'] = y
 

onnx2tf/ops/QuantizeLinear.py

@@ -11,6 +11,49 @@ from onnx2tf.utils.common_functions import (
     make_tf_node_info,
     convert_axis,
 )
+from onnx2tf.utils.enums import ONNX_DTYPES_TO_TF_DTYPES
+
+
+def _get_qmin_qmax(dtype: tf.dtypes.DType):
+    if dtype == tf.uint8:
+        return 0.0, 255.0
+    if dtype == tf.int8:
+        return -128.0, 127.0
+    if dtype == tf.uint16:
+        return 0.0, 65535.0
+    if dtype == tf.int16:
+        return -32768.0, 32767.0
+    return None, None
+
+
+def _expand_scale_or_zero_point(
+    *,
+    value,
+    input_tensor,
+    axis: int,
+    block_size: int,
+):
+    value_rank = len(value.shape)
+    input_rank = len(input_tensor.shape)
+
+    if value_rank == 0:
+        return value
+
+    if block_size > 0 and value_rank == input_rank:
+        if value.shape[axis] is None \
+            or input_tensor.shape[axis] is None \
+            or value.shape[axis] != input_tensor.shape[axis]:
+            expanded = tf.repeat(value, repeats=block_size, axis=axis)
+            expanded = tf.slice(expanded, [0] * input_rank, tf.shape(input_tensor))
+            return expanded
+        return value
+
+    if value_rank == 1 and input_rank is not None:
+        shape = [1] * input_rank
+        shape[axis] = -1
+        return tf.reshape(value, shape)
+
+    return value
 
 
 @print_node_info
@@ -60,12 +103,12 @@ def make_node(
 
     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
-    input_tensor_shape = input_tensor.shape
-    input_tensor_rank = len(input_tensor_shape)
+    input_nhwc = False
+    if isinstance(graph_node_input_1, gs.Variable):
+        input_nhwc = tf_layers_dict.get(graph_node_input_1.name, {}).get('nhwc', False)
+    input_tensor_rank = len(input_tensor.shape)
     y_scale = tf_layers_dict[graph_node_input_2.name]['tf_node'] \
         if isinstance(graph_node_input_2, gs.Variable) else graph_node_input_2
-    y_scale_shape = y_scale.shape
-    y_scale_rank = len(y_scale_shape)
     y_zero_point = tf_layers_dict[graph_node_input_3.name]['tf_node'] \
         if isinstance(graph_node_input_3, gs.Variable) else graph_node_input_3
 
@@ -81,6 +124,8 @@ def make_node(
         'optype': graph_node.op,
         'shape': shape,
         'dtype': dtype,
+        'is_dequantized': True,
+        'nhwc': input_nhwc,
     }
 
     # Generation of TF OP
@@ -88,51 +133,79 @@ def make_node(
         x=input_tensor,
         dtype=tf.float32,
     )
-    x_shape = input_tensor_shape
-    x_rank = input_tensor_rank
-    y_scale_shape = y_scale_shape
-
-    # Reshape process is needed for per-axis quantization
-    # when scale is a 1-D tensor
-    if y_scale_rank == 1:
-        shape_broadcast = list(
-            [1 for _ in range(axis)] \
-            + [x_shape[axis]] \
-            + [1 for _ in range(axis + 1, x_rank)]
-        )
-        y_scale = tf.reshape(
-            tensor=y_scale,
-            shape=shape_broadcast,
-        )
-    y = tf.divide(
-        x=input_tensor,
-        y=y_scale,
+
+    # If QuantizeLinear is immediately followed by Cast -> DequantizeLinear
+    # or DequantizeLinear only, bypass fake-quant to avoid generating
+    # Mul/Round/Min/Relu/Mul chains in TF/TFLite.
+    bypass_fake_quant = False
+    if graph_node.outputs and len(graph_node.outputs) > 0:
+        consumers = graph_node.outputs[0].outputs
+        if consumers:
+            bypass_fake_quant = True
+            for consumer in consumers:
+                if consumer.op == 'DequantizeLinear':
+                    continue
+                if consumer.op == 'Cast':
+                    cast_outs = consumer.outputs[0].outputs if consumer.outputs else []
+                    if not cast_outs or any(grand.op != 'DequantizeLinear' for grand in cast_outs):
+                        bypass_fake_quant = False
+                        break
+                else:
+                    bypass_fake_quant = False
+                    break
+
+    if bypass_fake_quant:
+        tf_layers_dict[graph_node_output.name]['tf_node'] = input_tensor
+        tf_layers_dict[graph_node_output.name]['tf_node_info'] = \
+            make_tf_node_info(
+                node_info={
+                    'tf_op_type': 'QuantizeLinear',
+                    'tf_inputs': {
+                        'x': input_tensor,
+                    },
+                    'tf_outputs': {
+                        'output': tf_layers_dict[graph_node_output.name]['tf_node'],
+                    },
+                }
+            )
+        return
+    y_scale = tf.cast(y_scale, tf.float32)
+
+    block_size = int(graph_node.attrs.get('block_size', 0))
+    y_scale = _expand_scale_or_zero_point(
+        value=y_scale,
+        input_tensor=input_tensor,
+        axis=axis,
+        block_size=block_size,
     )
-    y = tf.round(y)
 
-    if y_zero_point is not None:
-        y_dtype = y_zero_point.dtype if y_zero_point.dtype not in [tf.int8, tf.uint8] else tf.float32
-        y_zero_point = tf.cast(
-            x=y_zero_point,
-            dtype=tf.float32,
-        )
-        y_zero_point = tf.reshape(
-            tensor=y_zero_point,
-            shape=shape_broadcast,
-        ) if y_scale_rank == 1 else y_zero_point
-        y = tf.add(
-            x=y,
-            y=y_zero_point,
+    output_dtype_attr = int(graph_node.attrs.get('output_dtype', 0))
+    if y_zero_point is None:
+        output_dtype = ONNX_DTYPES_TO_TF_DTYPES.get(output_dtype_attr, tf.uint8) \
+            if output_dtype_attr != 0 else tf.uint8
+        y_zero_point = tf.zeros_like(y_scale)
+    else:
+        output_dtype = y_zero_point.dtype
+        y_zero_point = tf.cast(y_zero_point, tf.float32)
+        y_zero_point = _expand_scale_or_zero_point(
+            value=y_zero_point,
+            input_tensor=input_tensor,
+            axis=axis,
+            block_size=block_size,
         )
-    else:  # y_zero_point default dtype = uint8
-        y_dtype = tf.uint8
 
-    # Generation of TF OP
+    y = tf.round(tf.divide(input_tensor, y_scale))
+    y = tf.add(y, y_zero_point)
+
+    qmin, qmax = _get_qmin_qmax(output_dtype)
+    if qmin is not None and qmax is not None:
+        y = tf.clip_by_value(y, qmin, qmax)
+
+    # dequantize to float32 output
     tf_layers_dict[graph_node_output.name]['tf_node'] = \
-        tf.saturate_cast(
-            value=y,
-            dtype=y_dtype,
-            name=graph_node.name,
+        tf.multiply(
+            x=tf.subtract(y, y_zero_point),
+            y=y_scale,
         )
 
     # Generation of Debug Info

onnx2tf/ops/Split.py

@@ -124,6 +124,32 @@ def make_node(
         **kwargs,
     )
 
+    def _infer_split_axis_runtime(input_tensor, sum_split, fallback_axis):
+        if sum_split is None:
+            return tf.cast(fallback_axis, tf.int32)
+        shape = tf.shape(input_tensor)
+        eq = tf.equal(shape, tf.cast(sum_split, tf.int32))
+        mask = tf.cast(eq, tf.int32)
+        count = tf.reduce_sum(mask)
+        axis_from = tf.argmax(mask, axis=0, output_type=tf.int32)
+        fallback_axis_tensor = tf.cast(fallback_axis, tf.int32)
+        is_single = tf.cast(tf.equal(count, 1), tf.int32)
+        return axis_from * is_single + fallback_axis_tensor * (1 - is_single)
+
+    axis_for_split = axis
+    sum_split = None
+    split_list = None
+    if isinstance(split, np.ndarray):
+        split_list = list(split)
+    elif isinstance(split, (list, tuple)):
+        split_list = list(split)
+    if split_list is not None and len(split_list) > 1:
+        if len(split_list) == sum([1 for dim in split_list if isinstance(dim, (np.int64, int))]):
+            sum_split = int(np.sum(split_list))
+            axis_dim = input_tensor_shape[axis] if axis < len(input_tensor_shape) else None
+            if axis_dim is None or (isinstance(axis_dim, int) and axis_dim != sum_split):
+                axis_for_split = _infer_split_axis_runtime(input_tensor, sum_split, axis)
+
     # Generation of TF OP
     splited_tensors = None
     if (
@@ -225,18 +251,17 @@ def make_node(
                 num=None,
                 name=graph_node.name,
             )
-    elif isinstance(split, np.ndarray) \
+    elif isinstance(split, (list, tuple, np.ndarray)) \
         and len(list(split)) > 1 \
-        and np.prod(split) != 1 \
-        and isinstance(input_tensor_shape[axis], int) \
-        and len(split) == sum([1 for dim in split if isinstance(dim, np.int64) or isinstance(dim, int)]) \
-        and len(split) != sum([1 for dim in split if split[0] == dim]) \
-        and np.sum(split) == input_tensor_shape[axis]:
+        and (np.prod(split) != 1 if isinstance(split, np.ndarray) else True) \
+        and len(list(split)) == sum([1 for dim in list(split) if isinstance(dim, (np.int64, int))]) \
+        and len(list(split)) != sum([1 for dim in list(split) if list(split)[0] == dim]) \
+        and (not isinstance(input_tensor_shape[axis], int) or np.sum(list(split)) == input_tensor_shape[axis]):
         # Suppression of FlexSplitV generation
         # SplitV -> Strided_Slice
         splited_tensors = []
         begin_stock = []
-        for split_idx, split_dim in enumerate(split):
+        for split_idx, split_dim in enumerate(list(split)):
             begin_ = []
             end_ = []
             begin_mask_ = 0
@@ -269,7 +294,7 @@ def make_node(
             tf.split(
                 value=input_tensor,
                 num_or_size_splits=split,
-                axis=axis,
+                axis=axis_for_split,
                 num=num_outputs,
                 name=graph_node.name,
             )

{onnx2tf-1.29.15.dist-info → onnx2tf-1.29.17.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: onnx2tf
-Version: 1.29.15
+Version: 1.29.17
 Summary: Self-Created Tools to convert ONNX files (NCHW) to TensorFlow/TFLite/Keras format (NHWC). The purpose of this tool is to solve the massive Transpose extrapolation problem in onnx-tensorflow (onnx-tf).
 Keywords: onnx,tensorflow,tflite,keras,deep-learning,machine-learning
 Author: Katsuya Hyodo
@@ -364,7 +364,7 @@ Video speed is adjusted approximately 50 times slower than actual speed.
   docker run --rm -it \
   -v `pwd`:/workdir \
   -w /workdir \
-  ghcr.io/pinto0309/onnx2tf:1.29.15
+  ghcr.io/pinto0309/onnx2tf:1.29.17
 
   or
 
@@ -372,7 +372,7 @@ Video speed is adjusted approximately 50 times slower than actual speed.
   docker run --rm -it \
   -v `pwd`:/workdir \
   -w /workdir \
-  docker.io/pinto0309/onnx2tf:1.29.15
+  docker.io/pinto0309/onnx2tf:1.29.17
 
   or