Differentiation in arrays where len(X.shape)>2?

[yb6599]

Pertaining to this issue,

While trying to differentiate an array with len(X.shape) = 3, for example X.shape = (10, 100, 2) and t = np.arange(0, 10, 0.1), since X is not 2-dimensional, I get the error "Invalid Shape of X".
Without having to manipulate the shape of X, is it possible to implement differentiation of multidimensional arrays?

[andgoldschmidt]

It should be possible. The _align_axes(X, t, axis) and _restore_axes(dX, axis, flat) routines in differentiation.py can likely be adapted in a way that would enable the derived dimension to be any axis.

In the current version, the indices are arranged into batch x time, and the default behavior of derivative is to loop over each batch dimension. This doesn't leverage vectorization or multithreading, so it's not really performant. But, a possible minimal change would be to reshape internally into batch x time and then back into the original shape after the computation.