import torch


def _get_product_features(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """
    Get outer product of 2 tensors along the last dimension.
    All dimensions except last are preserved. The last dimension is replaced
        with flattened outer products of last-dimension-vectors from input tensors

    This is a vectorized implementation of (for 2D case):
    for i in range(x.shape[0]):
        out[i, :] = torch.outer(x[i, :], y[i, :]).flatten()

    For 2D inputs:
        Input shapes:
            x: (batch, feature_dim_x)
            y: (batch, feature_dim_y)
        Output shape:
            (batch, feature_dim_x*feature_dim_y)
    """
    return torch.einsum("...i,...j->...ij", (x, y)).flatten(start_dim=-2)