tico 0.1.0.dev250411__py3-none-any.whl

tico/serialize/circle_graph.py

@@ -0,0 +1,264 @@
+# 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 collections import defaultdict
+from typing import Any, cast, Dict, final, List, Optional, TYPE_CHECKING, Union
+
+if TYPE_CHECKING:
+    import torch.fx
+import numpy as np
+import torch
+from circle_schema import circle
+from torch._subclasses.fake_tensor import FakeTensor
+
+from tico.serialize.circle_mapping import (
+    extract_circle_dtype,
+    extract_shape,
+    str_to_circle_dtype,
+    to_circle_dtype,
+)
+from tico.serialize.pack import pack_buffer
+from tico.serialize.quant_param import QPARAM_KEY
+from tico.utils.utils import to_circle_qparam
+
+"""
+Type alias for const
+"""
+_PRIMITIVE_TYPES = (
+    float,
+    int,
+    bool,
+    str,
+    torch.Tensor,
+    torch.device,
+    torch.dtype,
+    torch.layout,
+)
+ConstDataElement = Union[
+    int, float, bool, str, torch.Tensor, torch.device, torch.dtype, torch.layout
+]
+ConstData = Union[ConstDataElement, List[ConstDataElement]]
+
+
+def is_const(arg) -> bool:
+    if isinstance(arg, FakeTensor):
+        return False
+    if isinstance(arg, _PRIMITIVE_TYPES):
+        return True
+    if isinstance(arg, (tuple, list)):
+        return all(map(is_const, arg))
+    if isinstance(arg, dict):
+        return all(map(is_const, arg.values()))
+    return False
+
+
+@final
+class CircleModel(circle.Model.ModelT):
+    def __init__(self):
+        super().__init__()
+        self.subgraphs: List[circle.SubGraph.SubGraphT] = []
+        self.buffers: List[circle.Buffer.BufferT] = []
+
+    def add_subgraph(self, graph: circle.SubGraph.SubGraphT) -> None:
+        self.subgraphs.append(graph)
+
+    def add_buffer(self, buffer: circle.Buffer.BufferT) -> int:
+        """Return buffer id"""
+        self.buffers.append(buffer)
+        buf_id = len(self.buffers) - 1  # last index
+        return buf_id
+
+
+@final
+class CircleSubgraph(circle.SubGraph.SubGraphT):
+    def __init__(self, model: CircleModel):
+        super().__init__()
+        self.model: CircleModel = model
+        self.name: str = "subgraph"
+        self.inputs: List[int] = []
+        self.outputs: List[int] = []
+        self.tensors: List[circle.Tensor.TensorT] = []
+        self.operators: List[circle.Operator.OperatorT] = []
+        self.name_to_tid: Dict[str, int] = {}
+        self.counter: defaultdict = defaultdict(int)
+
+    # Generate a unique name with prefix.
+    # Naming rule
+    # - If no tensor has the same name with prefix, return prefix
+    # - Otherwise, add postfix f"_{idx}" where idx increases by 1 from 0
+    # Example
+    # If prefix = "add", this function will find a unique name in the following order.
+    # "add", "add_0", "add_1", ...
+    def _gen_unique_name_with_prefix(self, prefix: str):
+        name = prefix
+        while self.has_tensor(name):
+            index = self.counter[prefix]
+            name = f"{prefix}_{index}"
+            self.counter[prefix] += 1
+
+        return name
+
+    def _add_tensor(self, tensor: circle.Tensor.TensorT) -> None:
+        self.tensors.append(tensor)
+        self.name_to_tid[tensor.name] = len(self.tensors) - 1
+
+    def add_operator(self, op: circle.Operator.OperatorT) -> None:
+        self.operators.append(op)
+
+    def add_input(self, input_name: str) -> None:
+        assert input_name in self.name_to_tid, f"{input_name}"
+        tid = self.name_to_tid[input_name]
+        self.inputs.append(tid)
+
+    def add_output(self, output: Any) -> None:
+        if isinstance(output, str):
+            assert output in self.name_to_tid
+            output_name = output
+        elif isinstance(output, int | float):
+            # output is built-in type.
+            circle_tensor = self.add_const_tensor(output)
+            output_name = circle_tensor.name
+        else:
+            raise NotImplementedError(f"Unsupported output dtype: {type(output)}")
+        tid = self.name_to_tid[output_name]
+        self.outputs.append(tid)
+
+    def has_tensor(self, name: str):
+        return name in self.name_to_tid
+
+    def add_tensor_from_node(
+        self, node: torch.fx.node.Node, data: Optional[np.ndarray] = None
+    ) -> None:
+        tensor = circle.Tensor.TensorT()
+        tensor.name = self._gen_unique_name_with_prefix(node.name)
+        assert node.meta.get("val") is not None
+        tensor.type = extract_circle_dtype(node)
+        tensor.shape = list(extract_shape(node))
+        if QPARAM_KEY in node.meta:
+            tensor.quantization = to_circle_qparam(node.meta[QPARAM_KEY])
+            tensor.type = str_to_circle_dtype(node.meta[QPARAM_KEY].dtype)
+
+        buffer = circle.Buffer.BufferT()
+        if data is not None and isinstance(data, np.ndarray):
+            data = data.flatten()
+
+            if QPARAM_KEY in node.meta:
+                if node.meta[QPARAM_KEY].dtype == "uint4":
+                    data = pack_buffer(data, "uint4")
+
+            # Packing np.ndarray is faster than packing bytes
+            buffer.data = data.view(np.uint8)  # type: ignore[assignment]
+        else:
+            assert data is None
+        bid = self.model.add_buffer(buffer)
+        tensor.buffer = bid
+        self._add_tensor(tensor)
+
+    def add_const_tensor(self, data: ConstData) -> circle.Tensor.TensorT:
+        assert is_const(data)
+        tensor = circle.Tensor.TensorT()
+        tensor.name = self._gen_unique_name_with_prefix("const_tensor")
+        assert not self.has_tensor(tensor.name)
+        torch_t = torch.as_tensor(data=data)
+        torch_t_shape = list(torch_t.size())
+        tensor.type = to_circle_dtype(torch_dtype=torch_t.dtype)
+        tensor.shape = torch_t_shape
+
+        buffer = circle.Buffer.BufferT()
+        buffer.data = torch_t.flatten().cpu().numpy().view(np.uint8)  # type: ignore[assignment]
+        bid = self.model.add_buffer(buffer)
+        tensor.buffer = bid
+        self._add_tensor(tensor)
+
+        return tensor
+
+    def add_tensor_from_scratch(
+        self, prefix: str, shape: List[int], dtype: int
+    ) -> circle.Tensor.TensorT:
+        assert isinstance(dtype, int), f"{dtype} must be integer. Use to_circle_dtype."
+        tensor = circle.Tensor.TensorT()
+        tensor.name = self._gen_unique_name_with_prefix(prefix)
+        tensor.type = dtype
+        tensor.shape = shape
+
+        buffer = circle.Buffer.BufferT()
+        bid = self.model.add_buffer(buffer)
+        tensor.buffer = bid
+        self._add_tensor(tensor)
+
+        return tensor
+
+    # Some operators like `full`, `arange_start_step` or `scalar_tensor` needs buffers to be in-place updated.
+    # TODO remove this function
+    def update_tensor_buffer(
+        self, data: ConstData, tensor_name: str = str()
+    ) -> circle.Tensor.TensorT:
+        assert is_const(data)
+        assert self.has_tensor(tensor_name)
+        data_tensor = torch.as_tensor(data=data)
+        data_shape = list(data_tensor.size())
+        op_tensor = self.tensors[self.name_to_tid[tensor_name]]
+        assert op_tensor.type == to_circle_dtype(
+            data_tensor.dtype
+        ), f"{op_tensor.type}, {data_tensor.dtype}"
+        assert op_tensor.shape == data_shape
+
+        buffer = circle.Buffer.BufferT()
+        # Packing np.ndarray is faster than packing bytes
+        buffer.data = data_tensor.flatten().cpu().numpy().view(np.uint8)  # type: ignore[assignment]
+        bid = self.model.add_buffer(buffer)
+        op_tensor.buffer = bid
+
+        return op_tensor
+
+    def get_tid_registered(
+        self, node: Union[torch.fx.node.Node, circle.Tensor.TensorT]
+    ) -> int:
+        assert hasattr(node, "name"), "FIX CALLER UNLESS"
+
+        tid = self.name_to_tid.get(node.name, None)
+
+        if tid is None:
+            raise KeyError(f"{node}({node.name}) is not registered.")
+
+        assert tid < len(self.tensors)
+
+        return tid
+
+    def get_tensor(self, node: torch.fx.node.Node) -> circle.Tensor.TensorT:
+        tid = self.get_tid_registered(node)
+
+        return self.tensors[tid]
+
+    def get_buffer(self, node: torch.fx.Node) -> circle.Buffer.BufferT:
+        buf_id = self.get_tensor(node).buffer
+        return self.model.buffers[buf_id]
+
+    # TODO Rename, it doesn't only get_tid but also possibly add a new const tensor
+    def get_tid(
+        self, node: Union[torch.fx.node.Node, circle.Tensor.TensorT, ConstData]
+    ) -> int:
+        # return -1 if node is None. This is for generating CircleOutputExclude
+        if node == None:
+            return -1
+
+        if hasattr(node, "name") and node.name in self.name_to_tid:
+            return self.name_to_tid[node.name]
+
+        if is_const(node):
+            node_name = self.add_const_tensor(cast(ConstData, node)).name
+            return self.name_to_tid[node_name]
+
+        # Unreachable
+        raise RuntimeError("fx Node was not converted to tensor.")

tico/serialize/circle_mapping.py

@@ -0,0 +1,177 @@
+# 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 Tuple, TYPE_CHECKING, Union
+
+if TYPE_CHECKING:
+    import torch.fx
+import numpy as np
+import torch
+from circle_schema import circle
+
+
+# Convert torch dtype to circle dtype
+def to_circle_dtype(
+    torch_dtype: torch.dtype,
+) -> int:
+    assert isinstance(torch_dtype, torch.dtype)
+    dmap = {
+        torch.float32: circle.TensorType.TensorType.FLOAT32,
+        torch.float: circle.TensorType.TensorType.FLOAT32,
+        torch.uint8: circle.TensorType.TensorType.UINT8,
+        torch.int8: circle.TensorType.TensorType.INT8,
+        torch.int16: circle.TensorType.TensorType.INT16,
+        torch.short: circle.TensorType.TensorType.INT16,
+        torch.int32: circle.TensorType.TensorType.INT32,
+        torch.int: circle.TensorType.TensorType.INT32,
+        torch.int64: circle.TensorType.TensorType.INT64,
+        torch.bool: circle.TensorType.TensorType.BOOL,
+    }
+
+    if torch_dtype not in dmap:
+        raise RuntimeError(f"Unsupported dtype {torch_dtype}")
+
+    circle_type = dmap[torch_dtype]
+    assert circle_type is not None
+    return circle_type
+
+
+# Convert str dtype used in QuantParam to circle dtype
+def str_to_circle_dtype(
+    str_dtype: str,
+) -> int:
+    dmap = {
+        "float32": circle.TensorType.TensorType.FLOAT32,
+        "float": circle.TensorType.TensorType.FLOAT32,
+        "uint8": circle.TensorType.TensorType.UINT8,
+        "int8": circle.TensorType.TensorType.INT8,
+        "int16": circle.TensorType.TensorType.INT16,
+        "short": circle.TensorType.TensorType.INT16,
+        "int32": circle.TensorType.TensorType.INT32,
+        "int": circle.TensorType.TensorType.INT32,
+        "int64": circle.TensorType.TensorType.INT64,
+        "bool": circle.TensorType.TensorType.BOOL,
+        "uint4": circle.TensorType.TensorType.UINT4,
+        # TODO Add more dtypes
+    }
+
+    if str_dtype not in dmap:
+        raise RuntimeError(f"Unsupported dtype {str_dtype}")
+
+    circle_type = dmap[str_dtype]
+    assert circle_type is not None
+    return circle_type
+
+
+# Convert circle dtype to numpy dtype
+def np_dtype_from_circle_dtype(circle_dtype: int):
+    dmap = {
+        circle.TensorType.TensorType.FLOAT32: np.float32,
+        circle.TensorType.TensorType.UINT8: np.uint8,
+        circle.TensorType.TensorType.INT8: np.int8,
+        circle.TensorType.TensorType.INT16: np.int16,
+        circle.TensorType.TensorType.INT32: np.int32,
+        circle.TensorType.TensorType.INT64: np.int64,
+        circle.TensorType.TensorType.BOOL: np.bool_,
+    }
+
+    if circle_dtype not in dmap:
+        raise RuntimeError(f"Unsupported dtype {circle_dtype}")
+
+    np_dtype = dmap[circle_dtype]
+    assert np_dtype is not None
+    return np_dtype
+
+
+# Return dtype of node
+def extract_torch_dtype(node: torch.fx.Node) -> torch.dtype:
+    assert node.meta is not None
+    assert node.meta.get("val") is not None
+
+    val = node.meta.get("val")
+    val_dtype = None
+    if isinstance(val, torch.Tensor):
+        assert isinstance(val.dtype, torch.dtype)
+        val_dtype = val.dtype
+    else:
+        val_dtype = torch.tensor(val).dtype
+    return val_dtype
+
+
+def extract_circle_dtype(node: torch.fx.Node) -> int:
+    return to_circle_dtype(extract_torch_dtype(node))
+
+
+# Return shape of node
+def extract_shape(node: torch.fx.Node) -> torch.Size:
+    assert node.meta is not None
+    assert node.meta.get("val") is not None
+
+    val = node.meta.get("val")
+    val_shape = None
+    if isinstance(val, torch.Tensor):
+        val_shape = val.size()
+    else:
+        val_shape = torch.tensor(val).shape
+
+    return val_shape
+
+
+# Return stride of node
+def extract_stride(node: torch.fx.Node) -> Tuple[int, ...]:
+    assert node.meta is not None
+    assert node.meta.get("val") is not None
+
+    val = node.meta.get("val")
+    val_stride = None
+    assert isinstance(val, torch.Tensor)
+    val_stride = val.stride()
+
+    return val_stride
+
+
+def traverse_elements(iter, container_types=(list, tuple)):
+    if isinstance(iter, container_types):
+        for e in iter:
+            for sub_e in traverse_elements(e, container_types):
+                yield sub_e
+    else:
+        yield iter
+
+
+def check_if_i32_range(axis: Union[list, int]):
+    INT32_MAX = 2**31 - 1
+    INT32_MIN = -(2**31)
+    values = list(traverse_elements(axis))
+    return all(INT32_MIN <= val <= INT32_MAX for val in values)
+
+
+def circle_legalize_dtype_to(values, *, dtype: torch.dtype):
+    """
+    Legalize data types from `torch.int64` to `torch.int32`.
+
+    Pytorch assumes python's built-in integer type is `torch.int64`.
+    But, many of the circle infrastructures support only int32 type. E.g. circle-interpreter.
+
+    So, if constants has values whose range is inside [INT32_MIN <= val <= INT32_MAX], we will legalize the data type to int32.
+
+    TODO support more types
+
+    NOTE. This function must be applied only to constant values.
+    """
+    if dtype != torch.int32:
+        raise RuntimeError("Not supported data types.")
+    if not check_if_i32_range(values):
+        raise RuntimeError("'size' cannot be converted from int64 to int32.")
+    return torch.as_tensor(values, dtype=dtype)

tico/serialize/circle_serializer.py

@@ -0,0 +1,232 @@
+# 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.
+
+import operator
+from typing import Dict
+
+import flatbuffers
+import torch
+from circle_schema import circle
+from torch.export.exported_program import (
+    ConstantArgument,
+    ExportedProgram,
+    InputKind,
+    TensorArgument,
+)
+
+from tico.serialize.circle_mapping import to_circle_dtype
+from tico.serialize.operators import *
+from tico.serialize.circle_graph import CircleModel, CircleSubgraph
+from tico.serialize.operators.hashable_opcode import OpCode
+from tico.serialize.operators.node_visitor import get_node_visitors
+from tico.utils import logging
+
+
+multiple_output_ops = [
+    torch.ops.aten.split_with_sizes.default,
+]
+
+# Build circle model from ExportedProgram
+# Return raw bytes of circle model
+def build_circle(edge_program: ExportedProgram) -> bytes:
+    logger = logging.getLogger(__name__)
+
+    builder = flatbuffers.Builder()
+
+    # Init Model
+    model = CircleModel()
+
+    # Add empty buffer at the front (convention)
+    model.add_buffer(circle.Buffer.BufferT())
+
+    # Create an empty subgraph (assume a single subgraph)
+    graph = CircleSubgraph(model)
+
+    # Export tensors
+    logger.debug("---------------Export tensors--------------")
+    buf_name_to_data = {name: buf for name, buf in edge_program.named_buffers()}
+    for node in edge_program.graph.nodes:
+        if node.op == "call_function":
+            if node.target in multiple_output_ops:
+                continue
+            node_val = node.meta["val"]
+            if node_val.layout != torch.strided:
+                raise RuntimeError(
+                    f"Only support dense tensors (node layout: {node_val.layout})"
+                )
+            graph.add_tensor_from_node(node)
+            logger.debug(f"call_function: {node.name} tensor exported.")
+
+        # placeholder: function input (including parameters, buffers, constant tensors)
+        elif node.op == "placeholder":
+            # placeholder invariants
+            assert node.args is None or len(node.args) == 0  # Not support default param
+
+            # parameters
+            if node.name in edge_program.graph_signature.inputs_to_parameters:
+                param_name = edge_program.graph_signature.inputs_to_parameters[
+                    node.name
+                ]
+                param_data = edge_program.state_dict[param_name]
+
+                assert isinstance(
+                    param_data, torch.Tensor
+                ), "Expect parameters to be a tensor"
+                param_value = param_data.cpu().detach().numpy()
+
+                graph.add_tensor_from_node(node, param_value)
+                logger.debug(f"placeholder(param): {node.name} tensor exported.")
+            elif node.name in edge_program.graph_signature.inputs_to_buffers:
+                buffer_name = edge_program.graph_signature.inputs_to_buffers[node.name]
+                assert buffer_name in buf_name_to_data
+                buffer_data = buf_name_to_data[buffer_name]
+                assert isinstance(
+                    buffer_data, torch.Tensor
+                ), "Expect buffers to be a tensor"
+                buffer_value = buffer_data.cpu().detach().numpy()
+
+                graph.add_tensor_from_node(node, buffer_value)
+                logger.debug(f"placeholder(buffer): {node.name} tensor exported.")
+            elif (
+                node.name
+                in edge_program.graph_signature.inputs_to_lifted_tensor_constants
+            ):
+                ctensor_name = (
+                    edge_program.graph_signature.inputs_to_lifted_tensor_constants[
+                        node.name
+                    ]
+                )
+                ctensor_data = edge_program.constants[ctensor_name]
+
+                assert isinstance(
+                    ctensor_data, torch.Tensor
+                ), "Expect constant tensor to be a tensor"
+                ctensor_value = ctensor_data.cpu().detach().numpy()
+
+                graph.add_tensor_from_node(node, ctensor_value)
+                logger.debug(
+                    f"placeholder(constant tensor): {node.name} tensor exported."
+                )
+            else:
+                user_inputs = [
+                    specs
+                    for specs in edge_program.graph_signature.input_specs
+                    if specs.kind == InputKind.USER_INPUT
+                ]
+                constant_inputs = [
+                    specs
+                    for specs in user_inputs
+                    if isinstance(specs.arg, ConstantArgument)
+                ]
+                name_to_value = {
+                    specs.arg.name: specs.arg.value for specs in constant_inputs
+                }
+                # NoneType ConstantArgument is ignored.
+                if node.name in name_to_value and name_to_value[node.name] == None:
+                    continue
+                graph.add_tensor_from_node(node)
+                logger.debug(f"placeholder: {node.name} tensor exported.")
+
+        # get_attr: retrieve parameter
+        elif node.op == "get_attr":
+            # node.name: Place where fetched attribute is saved
+            # node.target: Attribute in the module
+            attr_tensor = getattr(node.graph.owning_module, node.target)
+            assert isinstance(attr_tensor, torch.Tensor)
+
+            graph.add_tensor_from_scratch(
+                prefix=node.name,
+                shape=list(attr_tensor.shape),
+                dtype=to_circle_dtype(attr_tensor.dtype),
+            )
+
+            logger.debug(f"get_attr: {node.name} tensor exported.")
+
+        # output: function output
+        elif node.op == "output":
+            # output node itself does not need a buffer
+            # argument of output node is assumed to be exported beforehand
+            for output in node.args[0]:
+                if isinstance(output, torch.fx.Node):
+                    assert graph.has_tensor(output.name)
+            continue
+
+        # call_method: call method
+        elif node.op == "call_method":
+            raise AssertionError("Not yet implemented")
+
+        # call_module: call 'forward' of module
+        elif node.op == "call_module":
+            raise AssertionError("Not yet implemented")
+
+        else:
+            # Add more if fx.Node is extended
+            raise AssertionError(f"Unknown fx.Node op {node.op}")
+
+    # Register inputs
+    logger.debug("---------------Register inputs--------------")
+    for in_spec in edge_program.graph_signature.input_specs:
+        if in_spec.kind != InputKind.USER_INPUT:
+            continue
+        # NoneType ConstantArgument is ignored.
+        if isinstance(in_spec.arg, ConstantArgument) and in_spec.arg.value == None:
+            continue
+        arg_name = in_spec.arg.name
+        graph.add_input(arg_name)
+        logger.debug(f"Registered input: {arg_name}")
+
+    # Register outputs
+    logger.debug("---------------Register outputs--------------")
+    for user_output in edge_program.graph_signature.user_outputs:
+        graph.add_output(user_output)
+        logger.debug(f"Registered output: {user_output}")
+
+    # Export operators
+    logger.debug("---------------Export operators--------------")
+    op_codes: Dict[OpCode, int] = {}
+    visitors = get_node_visitors(op_codes, graph)
+    for node in edge_program.graph.nodes:
+        if node.op != "call_function":
+            continue
+
+        opcode = node.target
+        if opcode == operator.getitem:
+            continue
+        if opcode not in visitors:
+            raise RuntimeError(f"{opcode} is not yet supported")
+        circle_op = visitors[opcode].define_node(node)
+
+        if circle_op:
+            graph.add_operator(circle_op)
+            logger.debug(f"call_function: {node.name} ({opcode}) Op exported.")
+
+    # Register subgraph
+    model.subgraphs.append(graph)
+
+    # Encode operator codes
+    model.operatorCodes = [
+        code for code, _ in sorted(op_codes.items(), key=lambda x: x[1])
+    ]
+
+    # Description
+    model.description = "circle"
+
+    # Set version
+    model.version = 0
+
+    # Finish model
+    builder.Finish(model.Pack(builder), "CIR0".encode("utf8"))
+    buf = builder.Output()
+
+    return bytes(buf)

tico/serialize/operators/__init__.py

@@ -0,0 +1,28 @@
+# 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.
+
+import glob
+from os.path import basename, dirname, isfile, join
+
+from tico.utils.register_custom_op import RegisterOps
+
+
+# Register custom ops to torch namespace
+RegisterOps()
+
+# Load all modules in the current directory
+modules = glob.glob(join(dirname(__file__), "*.py"))
+__all__ = [
+    basename(f)[:-3] for f in modules if isfile(f) and not f.endswith("__init__.py")
+]