# mypy: allow-untyped-defs
r"""This file provides a location for operators that help exporting models via onnx.

E.g. `shape_as_tensor` and `reshape_from_tensor_shape`
are to make all dynamic sizes operations traceable.

NOTE: at one point these functions were implemented differently.
Since then we have implemented these directly in ATen, so this
file is kept purely for backward-compatibility.
"""

import torch
import torch.onnx


def shape_as_tensor(x):
    """Get the shape of a tensor as a tensor.

    Args:
        x (Tensor): The input tensor.

    Returns:
        Tensor: A tensor of shape [len(x.shape)] containing the size of each dimension of x.

    Example:
        >>> x = torch.randn(2, 3)
        >>> shape_as_tensor(x)
        tensor([2, 3])

    """
    return torch._shape_as_tensor(x)


def reshape_from_tensor_shape(x, shape):
    return torch._reshape_from_tensor(x, shape)