Function as_tensor

axlearn/common/utils.py:587–609

Converts `x` to Tensor recursively. Args: x: a jnp array, numpy array, TF/PyTorch Tensor, or a nested structure of arrays or Tensors. Returns: A nested structure with the same structure as `x` but with values converted to Tensors. Raises: NotImplementedError: I

(x: Any)

def as_tensor(x: Any):
    """Converts `x` to Tensor recursively.

    Args:
        x: a jnp array, numpy array, TF/PyTorch Tensor, or a nested structure of arrays or Tensors.

    Returns:
        A nested structure with the same structure as `x` but with values converted to Tensors.

    Raises:
        NotImplementedError: If conversion for the input type is unsupported.
    """
    if isinstance(x, Tensor):
        return x
    if isinstance(x, (numbers.Number, np.ndarray)):
        return jnp.asarray(x)
    if hasattr(x, "detach"):
        x = x.detach()
    if hasattr(x, "numpy"):
        return jnp.asarray(x.numpy())
    if isinstance(x, (Mapping, Sequence)):
        return jax.tree.map(as_tensor, x)
    raise NotImplementedError(f"{type(x)}: {x}")


def as_numpy_array(x: Any):