tico 0.1.0.dev250803__py3-none-any.whl → 0.1.0.dev251106__py3-none-any.whl

tico/serialize/operators/op_mm.py

@@ -20,7 +20,7 @@ if TYPE_CHECKING:
 import torch
 from circle_schema import circle
 
-from tico.serialize.circle_graph import CircleSubgraph, is_const
+from tico.serialize.circle_graph import CircleSubgraph
 from tico.serialize.operators.hashable_opcode import OpCode
 from tico.serialize.operators.node_visitor import NodeVisitor, register_node_visitor
 from tico.serialize.operators.utils import create_builtin_operator, get_op_index
@@ -28,9 +28,9 @@ from tico.utils.validate_args_kwargs import MatmulArgs
 
 
 @register_node_visitor
-class MatmulDefaultVisitor(NodeVisitor):
+class MatmulVisitor(NodeVisitor):
     """
-    Convert matmul to equavalent BatchMatMul or FullyConnected with Transpose.
+    Convert matmul to Circle BatchMatMul
     """
 
     target: List[torch._ops.OpOverload] = [torch.ops.aten.mm.default]
@@ -38,131 +38,7 @@ class MatmulDefaultVisitor(NodeVisitor):
     def __init__(self, op_codes: Dict[OpCode, int], graph: CircleSubgraph):
         super().__init__(op_codes, graph)
 
-    # NOTE: Matmul is equivalent to Batch MatMul (batch=1)
-    def define_bmm_node(self, inputs, outputs) -> circle.Operator.OperatorT:
-        def set_bmm_option(operator):
-            operator.builtinOptionsType = (
-                circle.BuiltinOptions.BuiltinOptions.BatchMatMulOptions
-            )
-            option = circle.BatchMatMulOptions.BatchMatMulOptionsT()
-            option.adjointLhs, option.adjointRhs = False, False
-            option.asymmetricQuantizeInputs = False
-            operator.builtinOptions = option
-
-        op_index = get_op_index(
-            circle.BuiltinOperator.BuiltinOperator.BATCH_MATMUL, self._op_codes
-        )
-        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
-        set_bmm_option(operator)
-
-        return operator
-
-    def define_transpose_node(self, inputs, outputs) -> circle.Operator.OperatorT:
-        def set_transpose_option(operator):
-            operator.builtinOptionsType = (
-                circle.BuiltinOptions.BuiltinOptions.TransposeOptions
-            )
-            option = circle.TransposeOptions.TransposeOptionsT()
-            operator.builtinOptions = option
-
-        transpose_op_index = get_op_index(
-            circle.BuiltinOperator.BuiltinOperator.TRANSPOSE, self._op_codes
-        )
-        operator = create_builtin_operator(
-            self.graph, transpose_op_index, inputs, outputs
-        )
-        set_transpose_option(operator)
-        return operator
-
-    def define_fc_node(self, inputs, outputs) -> circle.Operator.OperatorT:
-        def set_fc_option(operator):
-            operator.builtinOptionsType = (
-                circle.BuiltinOptions.BuiltinOptions.FullyConnectedOptions
-            )
-            option = circle.FullyConnectedOptions.FullyConnectedOptionsT()
-
-            option.fusedActivationFunction = (
-                circle.ActivationFunctionType.ActivationFunctionType.NONE
-            )
-            option.weightsFormat = (
-                circle.FullyConnectedOptionsWeightsFormat.FullyConnectedOptionsWeightsFormat.DEFAULT
-            )
-            option.keepNumDims = False
-            option.asymmetricQuantizeInputs = False
-            option.quantizedBiasType = circle.TensorType.TensorType.FLOAT32
-
-            operator.builtinOptions = option
-
-        fc_op_index = get_op_index(
-            circle.BuiltinOperator.BuiltinOperator.FULLY_CONNECTED, self._op_codes
-        )
-        operator = create_builtin_operator(self.graph, fc_op_index, inputs, outputs)
-        set_fc_option(operator)
-        return operator
-
-    """
-    Define FullyConnnected with Tranpose operator.
-    Note that those sets of operators are equivalent.
-    (1) Matmul
-    matmul( lhs[H, K], rhs[K, W'] ) -> output(H, W')
-    
-    (2) Transpose + FullyConneccted
-    transpose( rhs[K, W'] ) -> trs_output[W', K]
-    fullyconnected( lhs[H, K], trs_output[W', K] ) -> output(H, W')
-    """
-
-    def define_fc_with_transpose(
-        self, node, inputs, outputs
-    ) -> circle.Operator.OperatorT:
-        lhs, rhs = inputs
-
-        # get transpose shape
-        rhs_tid: int = self.graph.get_tid_registered(rhs)
-        rhs_tensor: circle.Tensor.TensorT = self.graph.tensors[rhs_tid]
-        rhs_name: str = rhs.name
-        rhs_type: int = rhs_tensor.type
-        rhs_shape: List[int] = rhs_tensor.shape
-        assert len(rhs_shape) == 2, len(rhs_shape)
-        rhs_shape_transpose = [rhs_shape[1], rhs_shape[0]]
-
-        # create transpose output tensor
-        trs_output = self.graph.add_tensor_from_scratch(
-            prefix=f"{rhs_name}_transposed_output",
-            shape=rhs_shape_transpose,
-            shape_signature=None,
-            dtype=rhs_type,
-            source_node=node,
-        )
-        trs_perm = self.graph.add_const_tensor(data=[1, 0], source_node=node)
-        trs_operator = self.define_transpose_node([rhs, trs_perm], [trs_output])
-        self.graph.add_operator(trs_operator)
-
-        # define fc node
-        fc_input = lhs
-        fc_weight = trs_output
-        fc_shape = [fc_weight.shape[0]]
-        fc_bias = self.graph.add_const_tensor(
-            data=[0.0] * fc_shape[0], source_node=node
-        )
-
-        operator = self.define_fc_node([fc_input, fc_weight, fc_bias], outputs)
-
-        return operator
-
-    def define_node(
-        self, node: torch.fx.Node, prior_latency=True
-    ) -> circle.Operator.OperatorT:
-        """
-        NOTE: Possibility of accuracy-latency trade-off
-        From ONE compiler's perspective:
-        - BMM uses per-tensor quantization for both rhs and lhs.
-        - FC uses per-channel quantization for weight and per-tensor for input.
-        Thus, FC is better in terms of accuracy.
-        FC necessarily involves an additional transpose operation to be identical with mm.
-        If transposed operand is const, it can be optimized by constant folding.
-        Thus, convert FC only if tranpose can be folded.
-        TODO set prior_latency outside
-        """
+    def define_node(self, node: torch.fx.Node) -> circle.Operator.OperatorT:
         args = MatmulArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
         input = args.input
         other = args.other
@@ -170,9 +46,16 @@ class MatmulDefaultVisitor(NodeVisitor):
         inputs = [input, other]
         outputs = [node]
 
-        if not is_const(other) and prior_latency:
-            operator = self.define_bmm_node(inputs, outputs)
-        else:
-            operator = self.define_fc_with_transpose(node, inputs, outputs)
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.BATCH_MATMUL, self._op_codes
+        )
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+        operator.builtinOptionsType = (
+            circle.BuiltinOptions.BuiltinOptions.BatchMatMulOptions
+        )
+        option = circle.BatchMatMulOptions.BatchMatMulOptionsT()
+        option.adjointLhs, option.adjointRhs = False, False
+        option.asymmetricQuantizeInputs = False
+        operator.builtinOptions = option
 
         return operator

tico/serialize/operators/op_rmsnorm.py

@@ -0,0 +1,65 @@
+# Copyright (c) 2025 Samsung Electronics Co., Ltd. All Rights Reserved
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Dict, List, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch._ops
+    import torch.fx
+import torch
+from circle_schema import circle
+
+from tico.serialize.circle_graph import CircleSubgraph
+from tico.serialize.operators.hashable_opcode import OpCode
+from tico.serialize.operators.node_visitor import NodeVisitor, register_node_visitor
+from tico.serialize.operators.utils import create_builtin_operator, get_op_index
+from tico.utils.validate_args_kwargs import CircleRMSNormArgs
+
+
+@register_node_visitor
+class RMSNormVisitor(NodeVisitor):
+    target: List[torch._ops.OpOverload] = [
+        torch.ops.circle_custom.rms_norm.default,
+    ]
+
+    def __init__(self, op_codes: Dict[OpCode, int], graph: CircleSubgraph):
+        super().__init__(op_codes, graph)
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+    ) -> circle.Operator.OperatorT:
+        args = CircleRMSNormArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        input = args.input
+        weight = args.weight
+        eps = args.eps
+
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.RMS_NORM, self._op_codes
+        )
+
+        inputs = [input, weight]
+        outputs = [node]
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        # Op-specific option
+        operator.builtinOptionsType = (
+            circle.BuiltinOptions.BuiltinOptions.RmsNormOptions
+        )
+        option = circle.RmsNormOptions.RmsNormOptionsT()
+        option.epsilon = eps
+
+        operator.builtinOptions = option
+
+        return operator

tico/utils/convert.py

@@ -20,26 +20,14 @@ import torch
 from torch.export import export, ExportedProgram
 
 from tico.config import CompileConfigBase, get_default_config
-from tico.experimental.quantization.passes.fold_quant_ops import FoldQuantOps
-from tico.experimental.quantization.passes.insert_quantize_on_dtype_mismatch import (
-    InsertQuantizeOnDtypeMismatch,
-)
-from tico.experimental.quantization.passes.propagate_qparam_backward import (
-    PropagateQParamBackward,
-)
-from tico.experimental.quantization.passes.propagate_qparam_forward import (
-    PropagateQParamForward,
-)
-from tico.experimental.quantization.passes.quantize_bias import QuantizeBias
-from tico.experimental.quantization.passes.remove_weight_dequant_op import (
-    RemoveWeightDequantOp,
-)
 from tico.passes.cast_aten_where_arg_type import CastATenWhereArgType
 from tico.passes.cast_clamp_mixed_type_args import CastClampMixedTypeArgs
 from tico.passes.cast_mixed_type_args import CastMixedTypeArgs
 from tico.passes.const_prop_pass import ConstPropPass
 from tico.passes.convert_conv1d_to_conv2d import ConvertConv1dToConv2d
+from tico.passes.convert_expand_to_slice_cat import ConvertExpandToSliceCat
 from tico.passes.convert_layout_op_to_reshape import ConvertLayoutOpToReshape
+from tico.passes.convert_matmul_to_linear import ConvertMatmulToLinear
 from tico.passes.convert_repeat_to_expand_copy import ConvertRepeatToExpandCopy
 from tico.passes.convert_to_relu6 import ConvertToReLU6
 from tico.passes.decompose_addmm import DecomposeAddmm
@@ -72,6 +60,14 @@ from tico.passes.remove_redundant_slice import RemoveRedundantSlice
 from tico.passes.remove_redundant_to_copy import RemoveRedundantToCopy
 from tico.passes.restore_linear import RestoreLinear
 from tico.passes.segment_index_select import SegmentIndexSelectConst
+from tico.quantization.passes.fold_quant_ops import FoldQuantOps
+from tico.quantization.passes.insert_quantize_on_dtype_mismatch import (
+    InsertQuantizeOnDtypeMismatch,
+)
+from tico.quantization.passes.propagate_qparam_backward import PropagateQParamBackward
+from tico.quantization.passes.propagate_qparam_forward import PropagateQParamForward
+from tico.quantization.passes.quantize_bias import QuantizeBias
+from tico.quantization.passes.remove_weight_dequant_op import RemoveWeightDequantOp
 from tico.serialize.circle_serializer import build_circle
 from tico.serialize.operators.node_visitor import get_support_targets
 from tico.utils import logging
@@ -141,6 +137,7 @@ def traced_run_decompositions(exported_program: ExportedProgram):
         or torch.__version__.startswith("2.7")
         or torch.__version__.startswith("2.8")
         or torch.__version__.startswith("2.9")
+        or torch.__version__.startswith("2.10")
     ):
         return run_decompositions(exported_program)
     else:
@@ -249,6 +246,14 @@ def convert_exported_module_to_circle(
             ConstPropPass(),
             SegmentIndexSelectConst(),
             LegalizeCausalMaskValue(enabled=config.get("legalize_causal_mask_value")),
+            ConvertExpandToSliceCat(enabled=config.get("convert_expand_to_slice_cat")),
+            ConvertMatmulToLinear(
+                enable_lhs_const=config.get("convert_lhs_const_mm_to_fc"),
+                enable_rhs_const=config.get("convert_rhs_const_mm_to_fc"),
+                enable_single_batch_lhs_const_bmm=config.get(
+                    "convert_single_batch_lhs_const_bmm_to_fc"
+                ),
+            ),
             LowerToResizeNearestNeighbor(),
             LegalizePreDefinedLayoutOperators(),
             LowerPow2ToMul(),
@@ -287,7 +292,7 @@ def convert_exported_module_to_circle(
 
     check_unsupported_target(exported_program)
     check_training_ops(exported_program)
-    circle_program = build_circle(exported_program)
+    circle_program = build_circle(exported_program, config)
 
     return circle_program
 

tico/utils/dtype.py

@@ -1,6 +1,8 @@
 import numpy as np
 import torch
 
+from circle_schema import circle
+
 NUMPY_TO_TORCH_DTYPE_DICT = {
     np.dtype("float32"): torch.float32,
     np.dtype("float64"): torch.float64,
@@ -15,6 +17,26 @@ NUMPY_TO_TORCH_DTYPE_DICT = {
     np.dtype("bool"): torch.bool,
 }
 
+CIRCLE_TO_TORCH_DTYPE_DICT = {
+    circle.TensorType.TensorType.FLOAT32: torch.float32,
+    circle.TensorType.TensorType.UINT8: torch.uint8,
+    circle.TensorType.TensorType.INT8: torch.int8,
+    circle.TensorType.TensorType.INT16: torch.int16,
+    circle.TensorType.TensorType.INT32: torch.int32,
+    circle.TensorType.TensorType.INT64: torch.int64,
+    circle.TensorType.TensorType.BOOL: torch.bool,
+}
+
 
 def numpy_dtype_to_torch_dtype(np_dtype: np.dtype) -> torch.dtype:
     return NUMPY_TO_TORCH_DTYPE_DICT[np_dtype]
+
+
+def circle_dtype_to_torch_dtype(circle_dtype: int) -> torch.dtype:
+    assert isinstance(circle_dtype, int)
+    if circle_dtype not in CIRCLE_TO_TORCH_DTYPE_DICT:
+        raise RuntimeError(f"Unsupported dtype {circle_dtype}")
+
+    torch_dtype = CIRCLE_TO_TORCH_DTYPE_DICT[circle_dtype]
+    assert torch_dtype is not None
+    return torch_dtype

tico/utils/register_custom_op.py

@@ -31,9 +31,11 @@ def CircleResizeNearestNeighbor():
         W_scale_factor = size[2] / W
         if H_scale_factor != W_scale_factor:
             raise RuntimeError("Scale factor of H and W should be same.")
-        return torch.nn.functional.interpolate(
-            input_, scale_factor=H_scale_factor, mode="nearest"
+        permuted = torch.permute(input_, [0, 3, 1, 2])
+        resized = torch.nn.functional.interpolate(
+            permuted, scale_factor=H_scale_factor, mode="nearest"
         )
+        return torch.permute(resized, [0, 2, 3, 1])
 
     @register_fake("circle_custom::resize_nearest_neighbor")
     def _(input_: torch.Tensor, size: List[int]):
@@ -631,7 +633,7 @@ def CircleInstanceNorm():
         bias: Optional[torch.Tensor] = None,
         running_mean: Optional[torch.Tensor] = None,
         running_var: Optional[torch.Tensor] = None,
-        use_input_stats: bool = False,
+        use_input_stats: bool = True,
         momentum: float = 0.1,
         eps: float = 1e-05,
         cudnn_enabled: bool = False,
@@ -639,7 +641,7 @@ def CircleInstanceNorm():
         NHWC_to_NCHW = [0, 3, 1, 2]
         NCHW_input = torch.ops.aten.permute.default(input_, NHWC_to_NCHW)
 
-        args = [NCHW_input, weight, bias, None, None, False, momentum, eps, False]
+        args = [NCHW_input, weight, bias, None, None, True, momentum, eps, False]
         NCHW_output = torch.ops.aten.instance_norm.default(*args)
         NCHW_to_NHWC = [0, 2, 3, 1]
         NHWC_output = torch.ops.aten.permute.default(NCHW_output, NCHW_to_NHWC)
@@ -703,6 +705,28 @@ def CircleQuantizeMX():
         return input_
 
 
+def CircleRMSNorm():
+    @custom_op("circle_custom::rms_norm", mutates_args=())
+    def rms_norm(
+        hidden_states: torch.Tensor,
+        weight: torch.Tensor,
+        eps: float = 1e-05,
+    ) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + eps)
+        return weight * hidden_states.to(input_dtype)
+
+    @register_fake("circle_custom::rms_norm")
+    def _(
+        hidden_states: torch.Tensor,
+        weight: torch.Tensor,
+        eps: float = 1e-05,
+    ) -> torch.Tensor:
+        return hidden_states.new_empty(hidden_states.size())
+
+
 # Add custom ops to the torch namespace
 def RegisterOps():
     CircleResizeNearestNeighbor()
@@ -715,3 +739,4 @@ def RegisterOps():
     CircleAvgPool2D()
     CircleInstanceNorm()
     CircleQuantizeMX()
+    CircleRMSNorm()

tico/utils/signature.py

@@ -0,0 +1,247 @@
+# Copyright (c) 2025 Samsung Electronics Co., Ltd. All Rights Reserved
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Sequence
+
+import numpy as np
+import torch
+from circle_schema import circle
+
+from tico.serialize.circle_mapping import to_circle_shape
+from tico.utils.dtype import circle_dtype_to_torch_dtype
+from tico.utils.installed_packages import is_dynamic_cache_available
+
+
+def is_dynamic_cache_instance(value):
+    if is_dynamic_cache_available():
+        from transformers.cache_utils import DynamicCache
+
+        return isinstance(value, DynamicCache)
+    else:
+        return False
+
+
+def flatten_and_convert_kwargs(kwargs: dict) -> dict[str, torch.Tensor]:
+    result = {}  # type: ignore[var-annotated]
+    for k, v in kwargs.items():
+        if v is None:
+            continue
+        elif isinstance(v, (list, tuple)):
+            # 1. handle list
+            def unpack_recursive(name, value, store=None):
+                if store is None:
+                    store = {}
+
+                if isinstance(value, (tuple, list)):
+                    for i, v in enumerate(value):
+                        # recursive call. Append index to name and explore lower level
+                        unpack_recursive(f"{name}_{i}", v, store)
+                else:
+                    # base type (scalar etc.) directly stored
+                    store[name] = value
+
+                return store
+
+            unpack_recursive(k, v, result)
+        elif is_dynamic_cache_instance(v):
+            # 2. handle DynamicCache
+            for idx, cache_val in enumerate(v.key_cache):
+                result[f"{k}_key_cache_{idx}"] = cache_val
+
+            for idx, cache_val in enumerate(v.value_cache):
+                result[f"{k}_value_cache_{idx}"] = cache_val
+        else:
+            result[k] = v
+
+    # 3. Convert to tensors
+    for k, v in result.items():
+        result[k] = v if isinstance(v, torch.Tensor) else torch.tensor(v)
+
+    return result
+
+
+def flatten_and_convert_args(args: Sequence) -> tuple:
+    result = []  # type: ignore[var-annotated]
+    for item in args:
+        if item is None:
+            continue
+
+        # 1. recursion on list and tuple
+        if isinstance(item, (list, tuple)):
+            result.extend(flatten_and_convert_args(item))
+            continue
+
+        # 2. handle DynamicCache
+        if is_dynamic_cache_available():
+            from transformers.cache_utils import DynamicCache
+
+            if isinstance(item, DynamicCache):
+                # NOTE The tensor order is: key_in → key_out → value_in → value_out
+                #
+                # Refer to https://github.com/huggingface/transformers/blob/3457e8e73e4f5532cc69059682b1ba4484d7e7e8/src/transformers/cache_utils.py#L557
+                # ```
+                # self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=-2)
+                # self.value_cache[layer_idx] = torch.cat([self.value_cache[layer_idx], value_states], dim=-2)
+                # ```
+                result.extend(item.key_cache)
+                result.extend(item.value_cache)
+                continue
+
+        # 3. Convert to tensors
+        result.append(item if isinstance(item, torch.Tensor) else torch.tensor(item))
+
+    return tuple(result)
+
+
+class ModelInputSpec:
+    @classmethod
+    def load(cls, circle_path):
+        def load(circle_path: str) -> bytes:
+            with open(circle_path, "rb") as f:
+                buf = bytes(f.read())
+            return buf
+
+        circle_binary = load(circle_path)
+        return cls(circle_binary)
+
+    def __init__(self, circle_binary):
+        model = circle.Model.Model.GetRootAsModel(circle_binary, 0)
+        assert model.SubgraphsLength() == 1, "Only one subgraph is supported"
+
+        graph = model.Subgraphs(0)
+        tensors = [graph.Tensors(graph.Inputs(o)) for o in range(graph.InputsLength())]
+
+        self.names = [t.Name().decode("utf-8").split("::")[-1] for t in tensors]
+        self.shapes = [t.ShapeAsNumpy() for t in tensors]
+        self.shape_signatures = list(
+            map(
+                lambda x: None if (isinstance(x, int) and x == 0) else x,
+                (t.ShapeSignatureAsNumpy() for t in tensors),
+            )
+        )
+        self.types: list[torch.dtype] = [
+            circle_dtype_to_torch_dtype(t.Type()) for t in tensors
+        ]
+        self.name_to_idx = {name: idx for idx, name in enumerate(self.names)}
+
+    def bind(self, args, kwargs, check=True):
+        """Convert args and kwargs into an ordered list according to model input order"""
+        inputs = []
+        args = flatten_and_convert_args(args)
+        kwargs = flatten_and_convert_kwargs(kwargs)
+
+        arg_num = len(args) + len(kwargs)
+        m_input_num = len(self.names)
+        if arg_num != m_input_num:
+            raise ValueError(
+                f"Mismatch: number of model inputs and number of passed arguments are not the same: inputs({m_input_num}) != passed({arg_num}), input spec: {self.names}"
+            )
+
+        # 1. positional arguments
+        for i, val in enumerate(args):
+            name = self.names[i]
+            inputs.append(val)
+
+        # 2. keyword arguments
+        for idx in range(len(args), len(self.names)):
+            name = self.names[idx]
+            inputs.append(kwargs[name])
+
+        if check:
+            self.check_types(inputs)
+            self.check_shapes(inputs)
+
+        return inputs
+
+    def check_types(self, inputs):
+        """Check the types of input values"""
+        for i, (inp, ref_type) in enumerate(zip(inputs, self.types)):
+            # TODO: Support more data types (np array)
+            assert isinstance(
+                inp, (torch.Tensor | int | float)
+            ), f"Input '{self.names[i]}' type must be a torch tensor or scalar."
+
+            if isinstance(inp, torch.Tensor):
+                if inp.dtype != ref_type:
+                    raise TypeError(
+                        f"Input '{self.names[i]}' type {inp.dtype} != expected {ref_type}"
+                    )
+            else:
+                # Scalars (int, float)
+                if ref_type == torch.float32:
+                    if not isinstance(inp, (float)):
+                        raise TypeError(
+                            f"Input '{self.names[i]}' type {type(inp)} != expected {ref_type}"
+                        )
+                elif ref_type == torch.int64:
+                    if not isinstance(inp, (int)):
+                        raise TypeError(
+                            f"Input '{self.names[i]}' type {type(inp)} != expected {ref_type}"
+                        )
+                else:
+                    print(f"Unexpected ref_type: {ref_type}")
+
+    def check_shapes(self, inputs):
+        """Check the shapes of input values"""
+
+        def merge(shape, shape_sig):
+            """
+            Merge shape signature with shape
+            """
+            from copy import deepcopy
+
+            shape_merged = deepcopy(shape)
+            if shape_sig is not None:
+                for idx, ss in enumerate(shape_sig):
+                    if ss == -1:
+                        shape_merged[idx] = -1
+
+            return shape_merged
+
+        for i, (inp, ref_shape, ref_shape_sig) in enumerate(
+            zip(inputs, self.shapes, self.shape_signatures)
+        ):
+            # TODO: Support more data types (np array)
+            assert isinstance(
+                inp, (torch.Tensor | int | float)
+            ), f"Input '{self.names[i]}' type must be a torch tensor or scalar."
+
+            if isinstance(inp, torch.Tensor):  # Tensor
+                in_shape, in_shape_sig = to_circle_shape(inp.size())
+
+                if len(in_shape) != len(ref_shape):
+                    raise ValueError(
+                        f"Input '{self.names[i]}' has invalid rank {len(in_shape)}!= expected {len(ref_shape)}"
+                    )
+
+                in_merged_shape = merge(in_shape, in_shape_sig)
+                ref_merged_shape = merge(ref_shape, ref_shape_sig)
+                for in_shp, ref_shp in zip(in_merged_shape, ref_merged_shape):
+                    if ref_shp == -1:
+                        continue
+                    if in_shp == -1:
+                        raise ValueError(
+                            f"Input '{self.names[i]}' has unknown dimension {inp.size()} != expected shape({ref_shape}) / shape signature({ref_shape_sig}) "
+                        )
+                    if in_shp != ref_shp:
+                        raise ValueError(
+                            f"Input '{self.names[i]}' has wrong dimension {inp.size()} != expected shape({ref_shape}) / shape signature({ref_shape_sig}) "
+                        )
+            elif isinstance(inp, (int, float)):  # Scalar
+                if len(ref_shape) > 0:
+                    raise ValueError(
+                        f"Input '{self.names[i]}' has invalid rank {len(ref_shape)}"
+                    )
+            else:
+                print(f"Unexpected input type: {type(inp)}")