Function cast

axlearn/common/utils.py:1164–1170

(x: Union[Tensor, TensorSpec])

            prefix = separator.join(path[: idx + 1]) if separator and is_str else path[: idx + 1]
            raise KeyError(f"No entries found at path '{prefix}'")
        x = x[key]

    return x


def set_recursively(
    x: NestedTensor,
    *,
    value: Tensor,
    path: Union[str, Sequence[str]],