def _expanded_shape(ndim, axis_size, axis):
    """
    Returns a shape with size = 1 for all dimensions
    except at axis.
    """
    return tuple([axis_size if i == axis else 1 for i in range(ndim)])