def slice_along_axis(index, axis, ndim):
    """
    Return a selector that takes a single-element slice (subtensor) of an nd-array
    along a certain axis (at a given index).
    The result would be an (n-1)-dimensional array. E.g. data[:, :, 5, :].
    The advantage of this function over subscript syntax is that you can
    specify the axis with a variable.

    :param index:
    :param axis:
    :param ndim:
    :return:
    """
    return tuple([index if axis == i else None for i in range(0, ndim)])