Function _eval_assign

python/tvm/script/parser/core/evaluator.py:563–604

Expression assignment evaluation implementation for TVMScript parser. Parameters ---------- target : doc.expr The root node of AST tree node of assigned expression to evaluate. source : Any The source to be assigned with evaluated expression. Returns ------

(
    target: doc.expr,
    source: Any,
)

def _eval_assign(
    target: doc.expr,
    source: Any,
) -> dict[str, Any]:
    """Expression assignment evaluation implementation for TVMScript parser.

    Parameters
    ----------
    target : doc.expr
        The root node of AST tree node of assigned expression to evaluate.

    source : Any
        The source to be assigned with evaluated expression.

    Returns
    -------
    res : Any
        The evaluation result.
    """
    target = doc.from_doc(target)
    assert isinstance(target, ast.expr)
    RHS_VAR_NAME = "__tvm_rhs_var__"  # pylint: disable=invalid-name
    rhs_var_name = RHS_VAR_NAME
    dict_locals = {rhs_var_name: source}
    mod = ast.fix_missing_locations(
        ast.Module(
            body=[
                ast.Assign(
                    targets=[target],
                    value=ast.Name(
                        id=rhs_var_name,
                        ctx=ast.Load(),
                    ),
                )
            ],
            type_ignores=[],
        )
    )
    exe = compile(mod, filename="<ast>", mode="exec")
    exec(exe, {}, dict_locals)  # pylint: disable=exec-used
    del dict_locals[rhs_var_name]
    return dict_locals