# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
"""The core infra for nn.Module, which includes the following pieces:
- Tensor, a wrapper on top of relax.Expr whose ty is a TensorType,
  providing more convenient access shape and dtype information.
  Tensor is always symbolic and not bound to any concrete values.
- Parameter, a special tensor which could be bound or not bound to concrete values.
- Module, a container of nn.Parameters and sub nn.Modules.
- Effect, a non-user-facing class that encloses potential side effects, for example, IO,
  impure external function callings, inplace mutation, etc.
"""

from collections import OrderedDict
from collections.abc import Callable, Iterator, Sequence
from typing import (
    TYPE_CHECKING,
    Any,
    Union,
)

import numpy as np  # type: ignore

import tvm.runtime
from tvm import tirx
from tvm.ir import IRModule
from tvm.ir.transform import Pass
from tvm.runtime import Device
from tvm.runtime import device as as_device
from tvm.runtime.vm import VirtualMachine
from tvm.target import Target

from .... import relax as rx
from ...block_builder import BlockBuilder
from ...type import (
    AnyType,
    ShapeType,
    TensorType,
    TupleType,
)
from ._tensor_op import _TensorOp
from .subroutine import SubroutineMixin

if TYPE_CHECKING:
    import torch  # type: ignore

    from . import spec as _spec
    from .extern import ExternModule


_DEFAULT_DTYPE = "float32"


def get_default_dtype() -> str:
    """Get the default parameter dtype if not specified. By default it is float32.

    Returns
    -------
    dtype : str
        The default dtype
    """
    return _DEFAULT_DTYPE


def set_default_dtype(dtype: str) -> None:
    """Set the default parameter dtype.

    Parameters
    ----------
    dtype : str
        The default dtype to be set
    """
    global _DEFAULT_DTYPE  # pylint: disable=global-statement
    _DEFAULT_DTYPE = dtype


class Tensor(_TensorOp):
    """A wrapper on top of relax.Expr whose ty is a TensorType, providing more
    convenient access shape and dtype information. Tensor is always symbolc and not bound to any
    concrete values. Shape and dtype inference is done eagerly upon tensor creation, i.e. when
    operators are applied on tensors, the shape and dtype information is already available.
    """

    _expr: rx.Expr

    def __init__(self, *, _expr: rx.Expr) -> None:
        """Private constructor. Tensor is never supposed to be constructed directly by users."""

        def _check_tensor(expr: rx.Expr) -> None:
            assert expr.ty is not None
            assert isinstance(expr.ty, TensorType)
            assert expr.ty.ndim != -1
            assert expr.ty.shape is not None
            assert expr.ty.shape.ty is not None
            assert isinstance(expr.ty.shape.ty, ShapeType)
            assert expr.ty.shape.ty.values is not None

        _check_tensor(_expr)
        self._expr = _expr

    @staticmethod
    def from_const(data) -> "Tensor":
        """Construct a tensor from numpy constants."""
        return Tensor(_expr=rx.const(data))

    @staticmethod
    def from_scalar(data: int | float, dtype: str) -> "Tensor":
        """Construct a tensor from a scalar with dtype specified."""
        return Tensor(_expr=rx.const(data, dtype=dtype))

    @staticmethod
    def from_ty(ty: rx.TensorType, name: str = "tensor") -> "Tensor":
        """Construct a nn.Tensor from a Relax TensorType.

        TensorType is the Relax type-level description of a tensor, carrying its shape
        and dtype without holding actual data. This factory creates an unbound placeholder
        ``nn.Tensor`` that can be used as a symbolic input when tracing an ``nn.Module``.

        Parameters
        ----------
        ty : rx.TensorType
            The type describing the tensor's shape and dtype.

        name : str
            Name hint for the underlying Relax variable.

        Returns
        -------
        tensor : Tensor
            A symbolic ``nn.Tensor`` backed by a ``relax.Var`` with the given type.
        """
        return Tensor(
            _expr=rx.Var(
                name_hint=name,
                ty=ty,
            )
        )

    @staticmethod
    def placeholder(
        shape: Sequence[int | str | tirx.Expr],
        dtype: str,
        name: str = "tensor",
    ) -> "Tensor":
        """Create a placeholder tensor with given shape and dtype. A placeholder tensor should
        never be created directly by users in usual cases, and the only exception is to indicate
        the shape/dtype of return values of an external function.

        If shape is a string `name`, we create a symbolic shape `tvm.tirx.Var(name, "int64")`.
        """
        new_shape = []
        for expr in shape:
            if isinstance(expr, int | tirx.IntImm):
                expr = int(expr)
                assert expr >= 0
                new_shape.append(expr)
                continue
            if isinstance(expr, str):
                expr = tirx.Var(expr, "int64")
                new_shape.append(expr)
                continue
            if not tvm.ir.is_prim_expr(expr):
                raise TypeError(f"Invalid shape: {shape}")
            assert expr.ty == tvm.ir.PrimType("int64")
            new_shape.append(expr)
        return Tensor(
            _expr=rx.Var(
                name_hint=name,
                ty=TensorType(
                    shape=new_shape,  # type: ignore[arg-type]
                    dtype=dtype,
                ),
            )
        )

    @property
    def shape(self) -> list[int | tirx.Expr]:
        """Returns the shape of the tensor as a list of integers.

        An integer can be a python int or tvm.tirx.Expr, depending on whether the shape is
        fully static, for example, [1, 2, tvm.tirx.Var("n")] is a valid shape where the last
        dimension is dynamic while the first two dimensions are always static constants.

        Returns
        -------
        shape : List[Union[int, tirx.Expr]]
            The shape of the tensor
        """

        def _simplify(expr: tirx.Expr):
            return expr.value if isinstance(expr, tirx.IntImm) else expr

        shape_ty: ShapeType = self._expr.ty.shape.ty
        return [_simplify(x) for x in shape_ty.values]

    @property
    def ndim(self) -> int:
        """Returns the number of dimensions of the tensor.

        Returns
        -------
        ndim : int
            The number of dimensions of the tensor
        """
        return self._expr.ty.ndim

    @property
    def dtype(self) -> str:
        """Returns the data type of the tensor.

        Returns
        -------
        dtype : str
            The data type of the tensor
        """
        return self._expr.ty.dtype

    def __repr__(self) -> str:
        return f'Tensor({self.shape}, "{self.dtype}")'


class Parameter(Tensor):
    """A parameter represents the weight of a neural network layer. It is a special tensor which
    could be bound or not bound to concrete values. If a parameter is bound to a concrete value,
    it is called a bound parameter, otherwise it is called an unbound parameter.
    """

    _data: Tensor | None
    attrs: dict[str, Any]

    def __init__(
        self,
        shape: Sequence[int | str | tirx.Expr],
        dtype: str | None = None,
    ) -> None:
        """Create a parameter with given shape and dtype. The parameter is not bound to any
        concrete values.

        Parameters
        ----------
        shape : Sequence[Union[int, str, tirx.Expr]]
            The shape of the parameter. If it is a string `name`, we create a symbolic shape
            `tvm.tirx.Var(name, "int64")`.
        dtype : Optional[str]
            The data type of the parameter. If not specified, the default dtype will be used.
        """
        if dtype is None:
            dtype = get_default_dtype()
        super().__init__(_expr=Tensor.placeholder(shape, dtype=dtype, name="param")._expr)
        self._data = None
        self.attrs = OrderedDict()

    @property
    def data(self) -> Tensor | None:
        """Returns the concrete value of the parameter if it is bound to a concrete value,
        otherwise returns None. The returned value is a tvm.runtime.Tensor."""
        return self._data

    @data.setter
    def data(self, data: Union[None, tvm.runtime.Tensor, np.ndarray, "torch.Tensor"]) -> None:
        """Set the concrete value of the parameter. The data should be one of the following:
        - None: unbind the parameter to concrete values
        - tvm.runtime.Tensor
        - numpy.ndarray
        - torch.Tensor and any other DLPack-compliant tensors
        """
        if data is None:
            self._data = data
            return
        # Try to do zero-copy if possible
        if isinstance(data, tvm.runtime.Tensor):
            pass
        elif isinstance(data, np.ndarray):
            data = tvm.runtime.tensor(data)
        elif hasattr(data, "__dlpack__"):
            data = _from_dlpack(data)
        else:
            raise TypeError(f"Unsupported data type: {type(data)}")
        if data.shape != tuple(self.shape):
            raise ValueError(f"Shape mismatch: expected {tuple(self.shape)}, got {data.shape}")
        if data.dtype != self.dtype:
            raise ValueError(f"Dtype mismatch: expected {self.dtype}, got {data.dtype}")
        self._data = data

    def to(self, dtype: str | None = None) -> None:  # pylint: disable=invalid-name
        """Change the dtype of the parameter if it is not bound to any concrete data"""
        if dtype is not None:
            if self._data is not None:
                raise ValueError(
                    "Changing the dtype of a Parameter that has been bound to concrete "
                    "data is not recommended. It might lead to potential precision loss "
                    "or other unexpected behaviors"
                )
            self._expr = Tensor.placeholder(  # pylint: disable=protected-access
                self.shape, dtype=dtype, name="param"
            )._expr


class Object:
    """A wrapper on top of relax.Expr whose ty is the base
    AnyType, rather than a more specific subtype. Object effectively
    represents non-tensor frontend components such as KV caches.
    """

    _expr: rx.Var

    def __init__(self, *, _expr: rx.Expr, _name: str) -> None:
        """Private constructor. Object is never supposed to be constructed directly by users."""
        if not isinstance(_expr, rx.Var):
            _expr = BlockBuilder.current().emit(_expr, _name)
        self._expr = _expr
        assert isinstance(self._expr.ty, AnyType)


class Effect:
    """Effect is a special non-user facing type that is used to represent operations with side
    effects, for example, print. It is used to represent the output of a computation.
    """

    def emit_init(self, name_hint: str, builder: BlockBuilder) -> list[rx.DataflowVar]:
        """Emit the initialization of the effect. This method is called by the compiler to
        initialize the effect."""
        raise NotImplementedError

    def create(self, name_hint: str) -> list[rx.Var]:
        """Create the implicit inputs to a relax.Function that represents the side effect"""
        raise NotImplementedError

    def set_state(self, state_vars: list[rx.Var]) -> None:
        """Set the variables that represents the effect"""
        raise NotImplementedError

    def finalize(self) -> list[rx.Var]:
        """finalize the effect as the implicit return value of a relax.Function"""
        raise NotImplementedError

    def to(self, dtype: str | None = None) -> None:  # pylint: disable=invalid-name
        """Convert the effect to specific dtype. Usually it is no-op for most of the effects"""


class Module(SubroutineMixin):
    """Base class for neural network components. Subclass it to build your models.
    Modules can nest within each other in a tree structure using regular attribute assignment."""

    def named_parameters(self, prefix: str = "") -> Iterator[tuple[str, Parameter]]:
        """This method provides an iterator over module parameters,
        yielding both the parameter name and its corresponding value.

        Parameters
        ----------
        prefix : str
            Prefix to prepend to all parameter names.

        Yields
        ------
        (str, Parameter) - Tuple containing the name and parameter
        """
        yield from _attribute_finder(
            self, prefix, condition_yield=lambda x: isinstance(x, Parameter)
        )

    def parameters(self) -> Iterator[Parameter]:
        """This method provides an iterator over module parameters,
        yielding only the Parameter value.

        Yields
        ------
        Parameter - The module's parameter
        """
        for _, param in self.named_parameters():
            yield param

    def state_dict(
        self, *, prefix: str = "", destination: dict[str, Parameter] | None = None
    ) -> dict[str, Parameter]:
        """Returns a dictionary containing references to the whole state of the module.

        Parameters
        ----------
        prefix : str
            Prefix to prepend to all parameter names.
        destination : Optional[Dict[str, Parameter]]
            Dictionary to which state will be saved. If None, a new dictionary is created.

        Returns
        -------
        dict : Dict[str, Parameter]
            a dictionary containing a whole state of the module
        """
        if destination is None:
            destination = OrderedDict()
        for name, param in _attribute_finder(
            self, prefix, condition_yield=lambda x: isinstance(x, Parameter)
        ):
            destination[name] = param
        return destination

    def load_state_dict(
        self, state_dict: dict[str, Parameter], strict: bool = True
    ) -> tuple[list[str], list[str]]:
        """This function copies parameters and buffers from the state_dict into the current module
        and its descendants. If `strict` is set to True, the keys in the `state_dict` must exactly
        match the keys returned by the `state_dict()` function of this module.

        Parameters
        ----------
        state_dict : Dict[str, Parameter]
            A dictionary containing a whole state of the module
        strict : bool = True
            Whether to strictly enforce that the keys in `state_dict` match the keys returned by
            this module's `state_dict()` function.

        Returns
        -------
        (missing_keys, unexpected_keys) : Tuple[List[str], List[str]]
            A tuple of two lists: the missing keys and the unexpected keys.
        """
        self_state_dict = self.state_dict()
        missing_keys: list[str] = []
        unexpected_keys: list[str] = []
        for key, value in state_dict.items():
            if key not in self_state_dict:
                unexpected_keys.append(key)
                continue
            if value.data is None:
                raise ValueError(f"Parameter {key} is not set to any concrete tensor")
            self_state_dict.pop(key).data = value.data
        missing_keys = list(self_state_dict.keys())
        if strict and (missing_keys or unexpected_keys):
            raise KeyError(f"Missing keys: {missing_keys}, Unexpected keys: {unexpected_keys}")
        return missing_keys, unexpected_keys

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        """Call the module with the given inputs and returns the output."""
        if not hasattr(self, "forward"):
            raise NotImplementedError(f"Module {type(self)} does not have a `forward` method")
        return self.forward(*args, **kwargs)  # pylint: disable=no-member

    def to(self, dtype: str | None = None) -> None:  # pylint: disable=invalid-name
        """Convert the module to specific dtype recursively"""
        for _, item in self.__dict__.items():
            if hasattr(item, "to") and callable(item.to):
                item.to(dtype=dtype)
        if dtype is not None and isinstance(getattr(self, "dtype", None), str):
            self.dtype = dtype  # pylint: disable=attribute-defined-outside-init

    def export_tvm(
        self,
        spec: "_spec.ModuleSpecType",
        debug: bool = False,
        allow_extern: bool = False,
    ) -> (
        tuple[IRModule, list[tuple[str, Parameter]]]
        | tuple[IRModule, list[tuple[str, Parameter]], list["ExternModule"]]
    ):
        """Export the module to TVM IRModule and parameters

        Parameters
        ----------
        spec : _spec.ModuleSpecType
            A dictionary mapping each input name to a specification
            that defines the inputs shape and dtype.
        debug : bool
            If set to True, then the exported module will support
            effects. This enables things like printing in the graph.

        Returns
        -------
        irmodule : tvm.ir.IRModule
            The converted tvm IR representation of the model.
        params : List[Tuple[str, Parameter]]
            A list of Parameters corresponding to the weights of the model.
        ext_mods : List[nn.ExternModule]
            A list of ExternModules that are used in the model.
        """
        # pylint: disable=import-outside-toplevel
        from . import spec as _spec
        from .exporter import Exporter

        # pylint: enable=import-outside-toplevel
        spec = _spec.ModuleSpec.from_raw(spec, self)
        mod, params, ext_mods = Exporter(debug=debug).build(spec)
        if allow_extern:
            return mod, params, ext_mods
        if ext_mods:
            raise ValueError(
                "`ExternModule`(s) exist when they are not allowed. "
                "Turn on flag `allow_extern` to allow."
            )
        return mod, params

    def jit(  # pylint: disable=too-many-arguments
        self,
        spec: "_spec.ModuleSpec",
        device: str | Device = "cpu",
        pipeline: None | str | Pass = "default_build",
        out_format: str = "torch",
        debug: bool = False,
    ) -> Any:
        """Just-in-time compile an ``nn.Module`` into a callable executable.

        The method exports the module to a Relax IRModule, applies the given compilation
        pipeline, builds a Relax VM executable, and wraps the result so it can be called
        directly (e.g. with PyTorch tensors when ``out_format="torch"``).

        Parameters
        ----------
        spec : _spec.ModuleSpec
            A specification mapping each module input to its shape and dtype.

        device : Union[str, Device]
            The device to compile and run on (e.g. ``"cpu"``, ``"cuda"``).

        pipeline : Union[None, str, Pass]
            The Relax compilation pipeline to apply. ``"default_build"`` uses the standard
            optimization pipeline; ``None`` skips pipeline passes.

        out_format : str
            Output wrapper format. ``"torch"`` returns a ``TorchModule`` whose ``forward``
            accepts and returns PyTorch tensors.

        debug : bool
            If ``True``, enable effect-based debugging (e.g. printing) in the compiled graph.

        Returns
        -------
        module : Any
            A callable wrapper (type depends on *out_format*) around the compiled VM.
        """

        def _compile(spec, device, pipeline, debug):
            # pylint: disable=import-outside-toplevel
            from ...transform import AttachExternModules
            from ...vm_build import build as relax_build
            from . import spec as _spec
            from .exporter import Exporter

            # pylint: enable=import-outside-toplevel

            spec = _spec.ModuleSpec.from_raw(spec, self)
            mod, params, ext_mods = Exporter(debug=debug).build(spec)
            mod = AttachExternModules(ext_mods)(mod)  # pylint: disable=not-callable
            vm = VirtualMachine(  # pylint: disable=invalid-name
                relax_build(
                    mod,
                    target=Target.from_device(device),
                    relax_pipeline=pipeline,
                ),
                device,
            )
            params = _param_to_tensor(params, device)
            return spec, vm, params

        device = as_device(device)
        spec, vm, params = _compile(spec, device, pipeline, debug)  # pylint: disable=invalid-name

        if out_format == "torch":
            from . import torch  # pylint: disable=import-outside-toplevel

            return torch.TorchModule(spec=spec, params=params, vm=vm)

        raise ValueError(f"Unknown out_format: {out_format}")


class ModuleDict(Module):
    """Holds submodules in a dict."""

    def __init__(self, modules: OrderedDict[str, Module] | None = None):
        if modules is None:
            self.modules = OrderedDict()
        else:
            self.modules = OrderedDict(modules)

    def __iter__(self):
        return iter(self.modules.values())

    def __getitem__(self, key: str) -> Module:
        return self.modules[key]

    def __setitem__(self, key: str, module: Module) -> None:
        self.modules[key] = module

    def __len__(self) -> int:
        return len(self.modules)

    def keys(self) -> Iterator[str]:
        return self.modules.keys()

    def values(self) -> Iterator[Module]:
        return self.modules.values()

    def items(self) -> Iterator[tuple[str, Module]]:
        return self.modules.items()

    def get(self, key: str, default: Module | None = None) -> Module | None:
        return self.modules.get(key, default)

    def update(self, modules: dict[str, Module]) -> None:
        self.modules.update(modules)

    def clear(self) -> None:
        self.modules.clear()

    def pop(self, key: str) -> Module:
        return self.modules.pop(key)

    def __contains__(self, key: str) -> bool:
        return key in self.modules

    def to(self, dtype: str | None = None) -> None:  # pylint: disable=invalid-name
        for module in self.modules.values():
            module.to(dtype=dtype)


class ParameterDict(Module):
    """Holds parameters in a dict."""

    def __init__(
        self,
        params: OrderedDict[str, Parameter] | dict[str, Parameter] | None = None,
    ):
        self.params: OrderedDict[str, Parameter] = OrderedDict()
        if params is not None:
            self.update(params)

    def __iter__(self) -> Iterator[str]:
        return iter(self.params)

    def __getitem__(self, key: str) -> Parameter:
        return self.params[key]

    def __setitem__(self, key: str, param: Parameter) -> None:
        if not isinstance(key, str):
            raise TypeError(f"ParameterDict keys must be strings, but got {type(key).__name__}")
        if not isinstance(param, Parameter):
            raise TypeError(
                f"ParameterDict values must be nn.Parameter, but got {type(param).__name__}"
            )
        self.params[key] = param

    def __len__(self) -> int:
        return len(self.params)

    def keys(self) -> Iterator[str]:
        return self.params.keys()

    def values(self) -> Iterator[Parameter]:
        return self.params.values()

    def items(self) -> Iterator[tuple[str, Parameter]]:
        return self.params.items()

    def get(self, key: str, default: Parameter | None = None) -> Parameter | None:
        return self.params.get(key, default)

    def update(self, params: dict[str, Parameter]) -> None:
        for key, param in params.items():
            self[key] = param

    def clear(self) -> None:
        self.params.clear()

    def pop(self, key: str) -> Parameter:
        return self.params.pop(key)

    def __contains__(self, key: str) -> bool:
        return key in self.params

    def to(self, dtype: str | None = None) -> None:  # pylint: disable=invalid-name
        for param in self.params.values():
            param.to(dtype=dtype)


class ModuleList(Module):
    """Holds submodules in a list."""

    def __init__(self, modules: list[Module]):
        self.modules = modules

    def __iter__(self):
        return iter(self.modules)

    def __getitem__(self, idx: int) -> Module:
        return self.modules[idx]

    def __setitem__(self, idx: int, module: Module) -> None:
        self.modules[idx] = module

    def __len__(self):
        return len(self.modules)

    def append(self, module: Module):
        """Add a module to the end of the ModuleList"""
        self.modules.append(module)

    def to(self, dtype: str | None = None) -> None:  # pylint: disable=invalid-name
        for module in self.modules:
            module.to(dtype=dtype)

    def forward(self, x):  # pylint: disable=invalid-name
        """Feed-forward pass of the module"""
        for module in self.modules:
            x = module(x)
        return x


class ParameterList(Module):
    """Holds parameters in a list."""

    def __init__(self, params: list[Parameter] | None = None):
        self.params: list[Parameter] = []
        if params is not None:
            self.extend(params)

    def __iter__(self) -> Iterator[Parameter]:
        return iter(self.params)

    def __getitem__(self, idx: int) -> Parameter:
        return self.params[idx]

    def __setitem__(self, idx: int, param: Parameter) -> None:
        if not isinstance(param, Parameter):
            raise TypeError(
                f"ParameterList elements must be nn.Parameter, but got {type(param).__name__}"
            )
        self.params[idx] = param

    def __len__(self) -> int:
        return len(self.params)

    def append(self, param: Parameter) -> None:
        """Add a parameter to the end of the ParameterList"""
        if not isinstance(param, Parameter):
            raise TypeError(
                f"ParameterList elements must be nn.Parameter, but got {type(param).__name__}"
            )
        self.params.append(param)

    def extend(self, params: list[Parameter]) -> None:
        """Add parameters to the end of the ParameterList"""
        for param in params:
            self.append(param)

    def to(self, dtype: str | None = None) -> None:  # pylint: disable=invalid-name
        for param in self.params:
            param.to(dtype=dtype)


def wrap_nested(expr: rx.Expr, name: str) -> Tensor | Sequence[Tensor]:
    """Wrap the given relax.Expr, emit it using the current BlockBuilder,
    and automatically handle nested cases if the expr represents a Tuple.

    Parameters
    ----------
    expr : relax.Expr
        The Expr to be wrapped.

    name : str
        Name hint.

    Returns
    -------
    result : Union[Tensor, Tuple[Tensor]]
        The computed result.
    """
    if not isinstance(expr, rx.DataflowVar):
        expr = BlockBuilder.current().emit(expr, name)
    if isinstance(expr.ty, TensorType):
        return Tensor(_expr=expr)
    if isinstance(expr.ty, TupleType):
        return tuple(
            wrap_nested(  # type: ignore
                rx.TupleGetItem(expr, i),
                name=f"{name}.{i}",
            )
            for i in range(len(expr.ty.fields))
        )
    raise TypeError(f"Unsupported return type: {expr.ty}")


def _attribute_finder(root: Module, prefix: str, condition_yield: Callable[[Any], bool]):
    """Find attributes that satisfy the condition recursively"""
    if isinstance(root, ParameterList):
        for i, param in enumerate(root):
            if condition_yield(param):
                yield prefix + f"{i}", param
        return
    elif isinstance(root, ParameterDict):
        for name, param in root.items():
            if condition_yield(param):
                yield prefix + name, param
        return
    elif isinstance(root, ModuleList):
        for i, subitem in enumerate(root):
            yield from _attribute_finder(subitem, prefix + f"{i}.", condition_yield)
        return
    elif isinstance(root, ModuleDict):
        for name, subitem in root.items():
            yield from _attribute_finder(subitem, prefix + f"{name}.", condition_yield)
        return
    for name, item in root.__dict__.items():
        if condition_yield(item):
            yield prefix + name, item
        elif isinstance(item, ParameterList):
            yield from _attribute_finder(
                item,
                prefix + name + ".",
                condition_yield,
            )
        elif isinstance(item, ParameterDict):
            yield from _attribute_finder(
                item,
                prefix + name + ".",
                condition_yield,
            )
        elif isinstance(item, ModuleList):
            yield from _attribute_finder(
                item,
                prefix + name + ".",
                condition_yield,
            )
        elif isinstance(item, ModuleDict):
            for sub_name, sub_item in item.items():
                yield from _attribute_finder(
                    sub_item,
                    prefix + name + f".{sub_name}.",
                    condition_yield,
                )
        elif isinstance(item, Module):
            yield from _attribute_finder(
                item,
                prefix + name + ".",
                condition_yield,
            )


def _from_dlpack(tensor) -> tvm.runtime.Tensor:
    try:
        return tvm.runtime.from_dlpack(tensor)
    except RuntimeError:
        pass
    # special logic for PyTorch
    device_type = tensor.device.type
    device_id = tensor.device.index or 0
    return tvm.runtime.tensor(
        tensor.numpy(),
        device=Device(
            Device._DEVICE_NAME_TO_TYPE[device_type],
            device_id,
        ),
    )


def _param_to_tensor(
    params: list[tuple[str, Parameter]], device: Device
) -> list[tvm.runtime.Tensor]:
    results = []
    missing = []
    for name, param in params:
        if param.data is None:
            missing.append(name)
        else:
            results.append(param.data.copyto(target=device))
    if missing:
        raise ValueError(f"Parameters are not set to any concrete values: {', '.join(missing)}")
    return results
