Function as_local_tensor

axlearn/common/test_utils.py:130–133

(x: Tensor)

def as_local_tensor(x: Tensor) -> NestedTensor:
    if isinstance(x, Tensor):
        return x
    raise NotImplementedError(f"{type(x)}: {x}")


def clean_hlo(hlo: str) -> str: