def index_axis(obj, i, a):
    """
    Return obj indexed at i on axis a
    This is done manually to support both lists and numpy arrays
    """
    if a ==  None:
        return obj
    elif a == 0:
        return obj[i]
    elif a == 1:
        return obj[:, i]
    else:
        raise NotImplementedError()