I would replace the current contiguous_axis and contiguous_batch_size with a version generalized to an arbitrary number of indices.
I’d be inclined to require the block sizes to be known at compile time, but if you’d prefer this not to be the case I would define them to return tuples of ArrayInterface.StaticInt for those known at compile time, and Int otherwise.
Let me know if you’d like to take a stab at it soon. Otherwise, I’ll get to it shortly.