Something like that should work. By default I think it would be on the first axis (axis=0) though, because scan goes over the first axis. You’d have to check the shapes.
Something like that should work. By default I think it would be on the first axis (axis=0) though, because scan goes over the first axis. You’d have to check the shapes.