onnx2tf 1.29.20__tar.gz → 1.29.22__tar.gz

{onnx2tf-1.29.20 → onnx2tf-1.29.22}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: onnx2tf
-Version: 1.29.20
+Version: 1.29.22
 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
@@ -18,7 +18,7 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Dist: requests==2.32.5
 Requires-Dist: numpy==1.26.4
-Requires-Dist: onnx==1.19.0
+Requires-Dist: onnx==1.19.1
 Requires-Dist: onnxruntime==1.23.0
 Requires-Dist: opencv-python==4.11.0.86
 Requires-Dist: onnxsim==0.4.36
@@ -365,7 +365,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.20
+  ghcr.io/pinto0309/onnx2tf:1.29.22
 
   or
 
@@ -373,18 +373,18 @@ 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.20
+  docker.io/pinto0309/onnx2tf:1.29.22
 
   or
 
-  pip install -U onnx==1.19.0 \
+  pip install -U onnx==1.19.1 \
   && pip install -U onnx-graphsurgeon==0.5.8 \
   && pip install -U onnxruntime==1.23.0 \
   && pip install -U onnxsim==0.4.36 \
   && pip install -U onnxoptimizer==0.4.2 \
   && pip install -U simple_onnx_processing_tools==1.1.32 \
-  && pip install -U sne4onnx>=1.0.13 \
-  && pip install -U sng4onnx>=1.0.4 \
+  && pip install -U sne4onnx==1.0.15 \
+  && pip install -U sng4onnx==1.0.5 \
   && pip install -U ai_edge_litert==1.2.0 \
   && pip install -U tensorflow==2.19.0 \
   && pip install -U protobuf==3.20.3 \
@@ -630,7 +630,7 @@ After many upgrades, the need for JSON parameter correction has become much less
 
 `-ois` an option to overwrite the input OP to a static size if it has undefined dimensions. `-cotof` option checks the accuracy of all OPs one by one. `-cotoa` is the error value of the threshold for determining an accuracy error. If there are undefined dimensions in the input OP, it is better to fix them to the static geometry to improve the accuracy of the accuracy measurement.
 
-Also, you can use the `-cind` option to specify custom input for `-cotof`, instead of using the default dummy input. Otherwise, all input values will be set to 1. For more information about the `-cind` option, please refer to [here](#cli-parameter).
+Also, you can use the `-cind` option to specify custom input for `-cotof`, instead of using the default dummy input. Otherwise, all input values will be set to 1. You can override the dummy input values with `--value_hints` (scalar only, `*:default` supported). For more information about the `-cind` option, please refer to [here](#cli-parameter).
 
 The `-cotof` option only compares the original ONNX and converted TensorFlow (Keras) models at Float32 precision, not at Float16 or INT8 precision.
 
@@ -644,6 +644,10 @@ onnx2tf -i mobilenetv2-12.onnx -b 1 -cotof -cotoa 1e-1
 or
 
 onnx2tf -i mobilenetv2-12.onnx -cotof -cotoa 1e-1 -cind "input" "/your/path/x.npy"
+
+or
+
+onnx2tf -i mobilenetv2-12.onnx -cotof -cotoa 1e-1 --value_hints "input:0.5" "*:1.0"
 ```
 ![image](https://user-images.githubusercontent.com/33194443/216901668-5fdb1e38-8670-46a4-b4b9-8a774fa7545e.png)
 
@@ -1826,6 +1830,14 @@ optional arguments:
     A value of 1 or more must be specified.
     Numerical values other than dynamic dimensions are ignored.
 
+  -vh VALUE_HINTS [VALUE_HINTS ...], \
+      --value_hints VALUE_HINTS [VALUE_HINTS ...]
+    Value hints for dummy inference input tensors.
+    The format is
+    "input_name_1:value" "input_name_2:value" "*:default_value"
+    "*" applies to all inputs not explicitly specified.
+    Values are scalar only.
+
   -nlt, --no_large_tensor
     Suppresses constant bloat caused by Tile OP when optimizing models in onnxsim.
     See: https://github.com/daquexian/onnx-simplifier/issues/178
@@ -2157,6 +2169,7 @@ convert(
   batch_size: Union[int, NoneType] = None,
   overwrite_input_shape: Union[List[str], NoneType] = None,
   shape_hints: Union[List[str], NoneType] = None,
+  value_hints: Union[List[str], NoneType] = None,
   no_large_tensor: Optional[bool] = False,
   output_nms_with_dynamic_tensor: Optional[bool] = False,
   switch_nms_version: Optional[str] = 'v4',
@@ -2377,6 +2390,13 @@ convert(
       A value of 1 or more must be specified.
       Numerical values other than dynamic dimensions are ignored.
 
+    value_hints: Optional[List[str]]
+      Value hints for dummy inference input tensors.
+      The format is
+      ['input_name_1:value', 'input_name_2:value', '*:default_value']
+      "*" applies to all inputs not explicitly specified.
+      Values are scalar only.
+
     no_large_tensor: Optional[bool]
       Suppresses constant bloat caused by Tile OP when optimizing models in onnxsim.
       See: https://github.com/daquexian/onnx-simplifier/issues/178
@@ -3032,6 +3052,7 @@ The above differences often cannot be dealt with by simply converting the model
 14. [nobuco](https://github.com/AlexanderLutsenko/nobuco)
 15. [onnx2torch](https://github.com/ENOT-AutoDL/onnx2torch)
 16. [ai-edge-torch](https://github.com/google-ai-edge/ai-edge-torch)
+17. [LiteRT.js](https://ai.google.dev/edge/litert/web)
 
 ## Acknowledgement
 1. https://github.com/onnx/models

{onnx2tf-1.29.20 → onnx2tf-1.29.22}/README.md

@@ -323,7 +323,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.20
+  ghcr.io/pinto0309/onnx2tf:1.29.22
 
   or
 
@@ -331,18 +331,18 @@ 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.20
+  docker.io/pinto0309/onnx2tf:1.29.22
 
   or
 
-  pip install -U onnx==1.19.0 \
+  pip install -U onnx==1.19.1 \
   && pip install -U onnx-graphsurgeon==0.5.8 \
   && pip install -U onnxruntime==1.23.0 \
   && pip install -U onnxsim==0.4.36 \
   && pip install -U onnxoptimizer==0.4.2 \
   && pip install -U simple_onnx_processing_tools==1.1.32 \
-  && pip install -U sne4onnx>=1.0.13 \
-  && pip install -U sng4onnx>=1.0.4 \
+  && pip install -U sne4onnx==1.0.15 \
+  && pip install -U sng4onnx==1.0.5 \
   && pip install -U ai_edge_litert==1.2.0 \
   && pip install -U tensorflow==2.19.0 \
   && pip install -U protobuf==3.20.3 \
@@ -588,7 +588,7 @@ After many upgrades, the need for JSON parameter correction has become much less
 
 `-ois` an option to overwrite the input OP to a static size if it has undefined dimensions. `-cotof` option checks the accuracy of all OPs one by one. `-cotoa` is the error value of the threshold for determining an accuracy error. If there are undefined dimensions in the input OP, it is better to fix them to the static geometry to improve the accuracy of the accuracy measurement.
 
-Also, you can use the `-cind` option to specify custom input for `-cotof`, instead of using the default dummy input. Otherwise, all input values will be set to 1. For more information about the `-cind` option, please refer to [here](#cli-parameter).
+Also, you can use the `-cind` option to specify custom input for `-cotof`, instead of using the default dummy input. Otherwise, all input values will be set to 1. You can override the dummy input values with `--value_hints` (scalar only, `*:default` supported). For more information about the `-cind` option, please refer to [here](#cli-parameter).
 
 The `-cotof` option only compares the original ONNX and converted TensorFlow (Keras) models at Float32 precision, not at Float16 or INT8 precision.
 
@@ -602,6 +602,10 @@ onnx2tf -i mobilenetv2-12.onnx -b 1 -cotof -cotoa 1e-1
 or
 
 onnx2tf -i mobilenetv2-12.onnx -cotof -cotoa 1e-1 -cind "input" "/your/path/x.npy"
+
+or
+
+onnx2tf -i mobilenetv2-12.onnx -cotof -cotoa 1e-1 --value_hints "input:0.5" "*:1.0"
 ```
 ![image](https://user-images.githubusercontent.com/33194443/216901668-5fdb1e38-8670-46a4-b4b9-8a774fa7545e.png)
 
@@ -1784,6 +1788,14 @@ optional arguments:
     A value of 1 or more must be specified.
     Numerical values other than dynamic dimensions are ignored.
 
+  -vh VALUE_HINTS [VALUE_HINTS ...], \
+      --value_hints VALUE_HINTS [VALUE_HINTS ...]
+    Value hints for dummy inference input tensors.
+    The format is
+    "input_name_1:value" "input_name_2:value" "*:default_value"
+    "*" applies to all inputs not explicitly specified.
+    Values are scalar only.
+
   -nlt, --no_large_tensor
     Suppresses constant bloat caused by Tile OP when optimizing models in onnxsim.
     See: https://github.com/daquexian/onnx-simplifier/issues/178
@@ -2115,6 +2127,7 @@ convert(
   batch_size: Union[int, NoneType] = None,
   overwrite_input_shape: Union[List[str], NoneType] = None,
   shape_hints: Union[List[str], NoneType] = None,
+  value_hints: Union[List[str], NoneType] = None,
   no_large_tensor: Optional[bool] = False,
   output_nms_with_dynamic_tensor: Optional[bool] = False,
   switch_nms_version: Optional[str] = 'v4',
@@ -2335,6 +2348,13 @@ convert(
       A value of 1 or more must be specified.
       Numerical values other than dynamic dimensions are ignored.
 
+    value_hints: Optional[List[str]]
+      Value hints for dummy inference input tensors.
+      The format is
+      ['input_name_1:value', 'input_name_2:value', '*:default_value']
+      "*" applies to all inputs not explicitly specified.
+      Values are scalar only.
+
     no_large_tensor: Optional[bool]
       Suppresses constant bloat caused by Tile OP when optimizing models in onnxsim.
       See: https://github.com/daquexian/onnx-simplifier/issues/178
@@ -2990,6 +3010,7 @@ The above differences often cannot be dealt with by simply converting the model
 14. [nobuco](https://github.com/AlexanderLutsenko/nobuco)
 15. [onnx2torch](https://github.com/ENOT-AutoDL/onnx2torch)
 16. [ai-edge-torch](https://github.com/google-ai-edge/ai-edge-torch)
+17. [LiteRT.js](https://ai.google.dev/edge/litert/web)
 
 ## Acknowledgement
 1. https://github.com/onnx/models

{onnx2tf-1.29.20 → onnx2tf-1.29.22}/onnx2tf/__init__.py

@@ -1,3 +1,3 @@
 from onnx2tf.onnx2tf import convert, main
 
-__version__ = '1.29.20'
+__version__ = '1.29.22'

{onnx2tf-1.29.20 → onnx2tf-1.29.22}/onnx2tf/onnx2tf.py

@@ -43,6 +43,7 @@ from typing import Optional, List, Any, Dict
 from argparse import ArgumentParser
 
 import importlib
+import onnx2tf.utils.common_functions as common_functions
 from onnx2tf.utils.common_functions import (
     dummy_onnx_inference,
     dummy_tf_inference,
@@ -639,6 +640,7 @@ def convert(
     batch_size: Optional[int] = None,
     overwrite_input_shape: Optional[List[str]] = None,
     shape_hints: Optional[List[str]] = None,
+    value_hints: Optional[List[str]] = None,
     no_large_tensor: Optional[bool] = False,
     output_nms_with_dynamic_tensor: Optional[bool] = False,
     switch_nms_version: Optional[str] = 'v4',
@@ -850,6 +852,15 @@ def convert(
         A value of 1 or more must be specified.\n
         Numerical values other than dynamic dimensions are ignored.
 
+    value_hints: Optional[List[str]]
+        Value hints for dummy inference input tensors.\n
+        The format is\n
+        ["input_name_1:value","input_name_2:value","*:default_value"].\n
+        "*" applies to all inputs not explicitly specified.\n
+        Values are scalar only.\n
+        e.g.\n
+        ['input0:0.5','mask:0','*:1.0']\n
+
     no_large_tensor: Optional[bool]
         Suppresses constant bloat caused by Tile OP when optimizing models in onnxsim.\n
         See: https://github.com/daquexian/onnx-simplifier/issues/178
@@ -1110,6 +1121,8 @@ def convert(
     if verbosity is None:
         verbosity = 'debug'
     set_log_level('error' if non_verbose else verbosity)
+    common_functions.set_dummy_shape_hints(shape_hints)
+    common_functions.set_dummy_value_hints(value_hints)
 
     # Either designation required
     if not input_onnx_file_path and not onnx_graph:
@@ -1326,6 +1339,10 @@ def convert(
                 'Failed to optimize the onnx file.'
             )
 
+    has_external_data = False
+    if input_onnx_file_path and os.path.exists(input_onnx_file_path):
+        has_external_data = check_has_external_data(input_onnx_file_path)
+
     # Automatic generation of each OP name - sng4onnx
     if not not_use_opname_auto_generate:
         info('')
@@ -1357,9 +1374,7 @@ def convert(
 
     # Loading Graphs
     # onnx_graph If specified, onnx_graph is processed first
-    has_external_data = False
     if not onnx_graph:
-        has_external_data = check_has_external_data(input_onnx_file_path)
         onnx_graph = onnx.load(input_onnx_file_path)
 
     if not auto_split_model and onnx_graph is not None:
@@ -1686,6 +1701,7 @@ def convert(
                     'batch_size': batch_size,
                     'overwrite_input_shape': overwrite_input_shape,
                     'shape_hints': shape_hints,
+                    'value_hints': value_hints,
                     'no_large_tensor': no_large_tensor,
                     'output_nms_with_dynamic_tensor': output_nms_with_dynamic_tensor,
                     'switch_nms_version': switch_nms_version,
@@ -3874,6 +3890,18 @@ def main():
             'Only applied to dynamic dimensions in inputs. \n' +
             'Only used when -cotof or -coto are specified.'
     )
+    parser.add_argument(
+        '-vh',
+        '--value_hints',
+        type=str,
+        nargs='+',
+        help=\
+            'Value hints for dummy inference input tensors. \n' +
+            'The format is\n' +
+            '"input_name_1:value" "input_name_2:value" "*:default_value". \n' +
+            '"*" applies to all inputs not explicitly specified. \n' +
+            'Values are scalar only.'
+    )
     parser.add_argument(
         '-nlt',
         '--no_large_tensor',
@@ -4359,6 +4387,7 @@ def main():
         batch_size=args.batch_size,
         overwrite_input_shape=args.overwrite_input_shape,
         shape_hints=args.shape_hints,
+        value_hints=args.value_hints,
         no_large_tensor=args.no_large_tensor,
         output_nms_with_dynamic_tensor=args.output_nms_with_dynamic_tensor,
         switch_nms_version=args.switch_nms_version,

{onnx2tf-1.29.20 → onnx2tf-1.29.22}/onnx2tf/ops/MaxPool.py

@@ -1,6 +1,7 @@
 import random
 random.seed(0)
 import numpy as np
+import itertools
 np.random.seed(0)
 import tensorflow as tf
 import tf_keras
@@ -119,6 +120,59 @@ def make_node(
                     **kwargs,
                 )
 
+    # Guard: brute-force axis alignment between NCHW and NHWC when batch dim mismatches.
+    # Only trigger when shapes are fully known to avoid destabilizing existing behavior.
+    def _shape_matches(shape_a, shape_b):
+        if shape_a is None or shape_b is None or len(shape_a) != len(shape_b):
+            return False
+        for dim_a, dim_b in zip(shape_a, shape_b):
+            if dim_a is None or dim_b is None:
+                continue
+            if dim_a != dim_b:
+                return False
+        return True
+
+    def _best_perm_to_match(cur_shape, target_shape):
+        rank = len(cur_shape)
+        best_perm = None
+        best_cost = None
+        for perm in itertools.permutations(range(rank)):
+            permuted = [cur_shape[i] for i in perm]
+            if not _shape_matches(permuted, target_shape):
+                continue
+            cost = sum(abs(i - perm[i]) for i in range(rank))
+            if best_cost is None or cost < best_cost:
+                best_cost = cost
+                best_perm = perm
+        return best_perm
+
+    try:
+        current_shape = input_tensor.shape.as_list()
+    except Exception:
+        current_shape = None
+
+    if onnx_input_shape is not None and current_shape is not None:
+        onnx_shape = [
+            dim if isinstance(dim, int) else None for dim in onnx_input_shape
+        ]
+        cur_shape = [
+            dim if isinstance(dim, int) else None for dim in current_shape
+        ]
+        if len(onnx_shape) in (3, 4, 5) \
+            and len(cur_shape) == len(onnx_shape) \
+            and None not in onnx_shape \
+            and None not in cur_shape:
+            expected_shape = [onnx_shape[0]] + onnx_shape[2:] + [onnx_shape[1]]
+            if cur_shape[0] != onnx_shape[0] \
+                and not _shape_matches(cur_shape, expected_shape):
+                perm = _best_perm_to_match(cur_shape, expected_shape)
+                if perm is not None:
+                    input_tensor = transpose_with_flexing_deterrence(
+                        input_tensor=input_tensor,
+                        perm=list(perm),
+                        **kwargs,
+                    )
+
     filter = None
 
     auto_pad = graph_node.attrs.get('auto_pad', 'NOTSET')

{onnx2tf-1.29.20 → onnx2tf-1.29.22}/onnx2tf/ops/ScatterElements.py

@@ -146,6 +146,48 @@ def make_node(
         axis=axis,
         indices=indices_tensor,
     )
+    indices_rank = None
+    if hasattr(indices_tensor, "shape") and indices_tensor.shape is not None:
+        try:
+            indices_rank = len(indices_tensor.shape)
+        except TypeError:
+            indices_rank = indices_tensor.shape.rank
+    updates_rank = updates_tensor_rank
+    broadcast_shape = None
+    pad_rank = 0
+
+    def _pad_and_broadcast(target_tensor, pad_rank, target_shape):
+        tensor = target_tensor
+        if isinstance(tensor, np.ndarray):
+            tensor = tf.convert_to_tensor(tensor)
+        if pad_rank <= 0:
+            return tf.broadcast_to(tensor, target_shape)
+        tensor_shape = tf.shape(tensor)
+        new_shape = tf.concat(
+            [tf.ones([pad_rank], dtype=tf.int32), tensor_shape],
+            axis=0,
+        )
+        tensor = tf.reshape(tensor, new_shape)
+        return tf.broadcast_to(tensor, target_shape)
+
+    updates_tensor_for_scatter = updates_tensor
+    if indices_rank is not None and updates_rank is not None and indices_rank != updates_rank:
+        if indices_rank > updates_rank:
+            broadcast_shape = tf.shape(indices_tensor)
+            pad_rank = indices_rank - updates_rank
+            updates_tensor_for_scatter = _pad_and_broadcast(
+                updates_tensor,
+                pad_rank,
+                broadcast_shape,
+            )
+        else:
+            broadcast_shape = tf.shape(updates_tensor)
+            pad_rank = updates_rank - indices_rank
+            indices_tensor = _pad_and_broadcast(
+                indices_tensor,
+                pad_rank,
+                broadcast_shape,
+            )
     sparsified_dense_idx_shape = updates_tensor_shape
 
     if None not in sparsified_dense_idx_shape:
@@ -160,6 +202,13 @@ def make_node(
         ]
 
     idx_tensors_per_axis = tf.meshgrid(*idx_tensors_per_axis, indexing='ij')
+    if indices_rank is not None \
+        and updates_rank is not None \
+        and indices_rank > updates_rank:
+        idx_tensors_per_axis = [
+            _pad_and_broadcast(idx_tensor, pad_rank, broadcast_shape)
+            for idx_tensor in idx_tensors_per_axis
+        ]
     idx_tensors_per_axis[axis] = indices_tensor
     dim_expanded_idx_tensors_per_axis = [
         tf.expand_dims(idx_tensor, axis=-1)
@@ -194,7 +243,7 @@ def make_node(
     )
 
     indices = tf.reshape(coordinate, [-1, input_tensor_rank])
-    updates = tf.reshape(updates_tensor, [-1])
+    updates = tf.reshape(updates_tensor_for_scatter, [-1])
     output = tf.tensor_scatter_nd_update(
         tensor=input_tensor,
         indices=indices,

{onnx2tf-1.29.20 → onnx2tf-1.29.22}/onnx2tf/ops/Unique.py

@@ -14,6 +14,7 @@ from onnx2tf.utils.common_functions import (
     make_tf_node_info,
     get_replacement_parameter,
     pre_process_transpose,
+    convert_axis,
 )
 from onnx2tf.utils.logging import Color
 
@@ -69,8 +70,20 @@ def make_node(
         **kwargs,
     )
 
+    input_tensor_shape = input_tensor.shape
+    tensor_rank = len(input_tensor_shape) \
+        if input_tensor_shape != tf.TensorShape(None) else 1
+
     axis = graph_node.attrs.get('axis', None)
     sorted = graph_node.attrs.get('sorted', 1)
+    if axis is not None:
+        if isinstance(axis, np.ndarray) and axis.shape == ():
+            axis = int(axis)
+        axis = convert_axis(
+            axis=int(axis),
+            tensor_rank=tensor_rank,
+            before_op_output_shape_trans=before_op_output_shape_trans,
+        )
 
     # Preserving Graph Structure (Dict)
     for graph_node_output in graph_node_outputs:
@@ -101,17 +114,64 @@ def make_node(
 
     # tf unique returns unsorted tensor, need to sort if option is enabled
     if sorted:
-        # TODO: implement sort
-        error_msg = f'' + \
-                    Color.RED(f'WARNING:') + ' ' + \
-                    f'Sort option in Unique ops is not implemented yet.'
-        print(error_msg)
-        assert False, error_msg
-
-    tf_layers_dict[graph_node_outputs[0].name]['tf_node'] = y
-    tf_layers_dict[graph_node_outputs[1].name]['tf_node'] = indices
-    tf_layers_dict[graph_node_outputs[2].name]['tf_node'] = inverse_indices
-    tf_layers_dict[graph_node_outputs[3].name]['tf_node'] = count
+        # Sort unique outputs to match ONNX sorted behavior.
+        def _argsort_supported(dtype):
+            return dtype.is_floating or dtype.is_integer or dtype == tf.bool
+
+        y_rank = y.shape.rank
+        axis_ = axis
+        if axis_ is None:
+            axis_ = 0
+        if axis_ < 0 and y_rank is not None:
+            axis_ = axis_ + y_rank
+
+        def _lexsort_perm(flat_2d):
+            if not _argsort_supported(flat_2d.dtype):
+                return None
+            cols = flat_2d.shape[1]
+            if cols is None:
+                return None
+            order = tf.range(tf.shape(flat_2d)[0])
+            for col in reversed(range(cols)):
+                col_vals = tf.gather(flat_2d, order)[:, col]
+                if col_vals.dtype == tf.bool:
+                    col_vals = tf.cast(col_vals, tf.int32)
+                order = tf.gather(order, tf.argsort(col_vals, stable=True))
+            return order
+
+        order = None
+        if y_rank is not None and y_rank == 1:
+            if _argsort_supported(y.dtype):
+                sort_vals = y
+                if sort_vals.dtype == tf.bool:
+                    sort_vals = tf.cast(sort_vals, tf.int32)
+                order = tf.argsort(sort_vals, stable=True)
+        elif y_rank is not None and axis_ is not None and 0 <= axis_ < y_rank:
+            perm = [axis_] + [i for i in range(y_rank) if i != axis_]
+            y_t = tf.transpose(y, perm)
+            flat = tf.reshape(y_t, [tf.shape(y_t)[0], -1])
+            order = _lexsort_perm(flat)
+
+        if order is None:
+            warn_msg = f'' + \
+                       Color.YELLOW(f'WARNING:') + ' ' + \
+                       f'Unique sort fallback to unsorted due to dynamic shape or unsupported dtype.'
+            print(warn_msg)
+        else:
+            y = tf.gather(y, order, axis=axis_)
+            count = tf.gather(count, order)
+            indices = tf.gather(indices, order)
+            inv_order = tf.argsort(order)
+            inverse_indices = tf.gather(inv_order, inverse_indices)
+
+    if len(graph_node_outputs) >= 1:
+        tf_layers_dict[graph_node_outputs[0].name]['tf_node'] = y
+    if len(graph_node_outputs) >= 2:
+        tf_layers_dict[graph_node_outputs[1].name]['tf_node'] = indices
+    if len(graph_node_outputs) >= 3:
+        tf_layers_dict[graph_node_outputs[2].name]['tf_node'] = inverse_indices
+    if len(graph_node_outputs) >= 4:
+        tf_layers_dict[graph_node_outputs[3].name]['tf_node'] = count
 
     # Generation of Debug Info
     tf_outputs = {f"output{idx}": value for idx, value in enumerate([y, indices, inverse_indices, count])}