def to_batch_atom_3_shape(positions):
    """Reshape a tensor or numpy array to have shape (batch_size, n_atoms, 3).

    This allows converting positions in flattened format, as yielded from
    the ``TrajectoryDataset`` object into the standard shape used by
    MDAnalysis trajectories.

    Parameters
    ----------
    positions : torch.Tensor, numpy.ndarray, or pint.Quantity
        The input can have the following shapes: (batch_size, n_atoms, 3),
        (batch_size, n_atoms * 3), (n_atoms, 3).

    Returns
    -------
    reshaped_positions : torch.Tensor, numpy.ndarray, or pint.Quantity
        A view of the original tensor or array with shape (batch_size, n_atoms, 3).

    """
    if (len(positions.shape) == 2 and positions.shape[-1] != 3 or
                len(positions.shape) == 3):
        # (batch_size, n_atoms * 3) or (batch_size, n_atoms, 3).
        batch_size = positions.shape[0]
    else:
        batch_size = 1

    if positions.shape[-1] != 3:
        n_atoms = positions.shape[-1] // 3
    else:
        n_atoms = positions.shape[-2]

    standard_shape = (batch_size, n_atoms, 3)
    if positions.shape != standard_shape:
        positions = positions.reshape(standard_shape)

    return positions