Hello,
I am writing a class handling NumPy arrays, and basing it off some computations I have done in C++ using the Eigen library. For tensors whose length are only known at runtime, the Eigen chip method is useful for slicing a tensor along a specific axis while keeping all other axes the same. I have not found a similar method in NumPy, but a simple implementation is
def chip(A, axis, vals):
return np.swapaxes(np.swapaxes(A, axis, -1)[..., vals], -1, axis)
Example usage would be (to slice the 3rd axis of a 4D array):
A = np.random.randn(10,11,12,13)
B = chip(A, axis=2, vals=np.arange(0,6))
B.shape #(10, 11, 6, 13)
Since this may be useful for others, despite its simplicity, I thought it may be useful to have something similar in NumPy.
Best,
Mike