def METHOD_NAME(obj):
    if isinstance(obj, jt.Var):
        return TensorWrapper(obj)
    elif isinstance(obj, dict):
        return {k: METHOD_NAME(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [METHOD_NAME(x) for x in obj]
    elif isinstance(obj, tuple):
        return tuple(METHOD_NAME(x) for x in obj)
    else:
        return obj