[Bug] `relax.transform.ToMixedPrecision` crashes on a local (closure) function inside a dataflow block

[YuanchengJiang]

The following code:

import tvm
from tvm import relax
from tvm.script import ir as I, relax as R

@I.ir_module
class Module:
    @R.function
    def main(x: R.Tensor((), "float32")):
        with R.dataflow():
            @R.function
            def bar(y: R.Tensor((), "float32")) -> R.Tensor((), "float32"):
                return R.add(y, y)
            gv = bar(x)
            R.output(gv)
        return gv

relax.analysis.well_formed(Module)  # passes
relax.transform.ToMixedPrecision()(Module)

Resulted in this output:

Traceback (most recent call last):
  File "min.py", line 18, in <module>
  File "src/relax/transform/to_mixed_precision.cc", line 620, in operator()
  File "src/relax/transform/to_mixed_precision.cc", line 612, in tvm::relax::ToMixedPrecision(tvm::relax::Function const&, DLDataType, tvm::ffi::Optional<tvm::ffi::Array<tvm::ffi::String, void>, void>)
  File "src/relax/transform/to_mixed_precision.cc", line 590, in tvm::relax::ToMixedPrecisionRewriter::VisitExpr_(tvm::relax::FunctionNode const*)
  File "src/relax/transform/to_mixed_precision.cc", line 576, in tvm::relax::ToMixedPrecisionRewriter::VisitBindingBlock_(tvm::relax::DataflowBlockNode const*)
  File "src/relax/transform/to_mixed_precision.cc", line 450, in tvm::relax::ToMixedPrecisionRewriter::VisitBinding(tvm::relax::Binding const&)
  File "src/relax/transform/to_mixed_precision.cc", line 444, in tvm::relax::ToMixedPrecisionRewriter::VisitExpr_(tvm::VarNode const*)
  File "src/relax/transform/to_mixed_precision.cc", line 435, in tvm::relax::ToMixedPrecisionRewriter::VisitVar_(tvm::Var const&)
  File "src/relax/transform/infer_amp_utils.cc", line 39, in tvm::relax::NTypeFrom(tvm::Expr const&, DLDataType)
  File "src/relax/transform/infer_amp_utils.cc", line 36, in tvm::relax::NTypeFrom(tvm::Type const&, DLDataType)
  File "src/relax/transform/infer_amp_utils.cc", line 28, in tvm::relax::NTypeFrom(const tvm::Type&, DLDataType)::<lambda(const tvm::Type&)>
tvm.error.InternalError: Check failed: (tensor) is false: Expected TensorType, but got R.Callable((R.Tensor((), dtype="float32"),), R.Tensor((), dtype="float32"), True)

To reproduce: python3 min.py

The module is well-formed. ToMixedPrecisionRewriter::VisitVar_ calls NTypeFrom(var) on every remapped var in a dataflow block. NTypeFrom accepts only (nested) TensorType, so a var bound to a local R.function (R.Callable) hits the ICHECK in infer_amp_utils.cc:28 instead of being skipped.

TVM commit: e269315c90e3a061c9e1c77b370ce883b1b223f4

This bug was found by fusion-fuzz

[Nanmur]

I reproduced this on the current main branch and submitted #20475 with a regression test. The fix keeps the dtype-restoration path for nested tensor variables, while preserving the remapped expression for non-tensor variables such as local functions. Thanks for the clear reproducer.