Batch Jacobian?

[jpreiss]

In our original use case of trajectory optimization where the NN is (a part of) the dynamics model, there are two stages where we need mlpfile:

1. roll out the dynamics around the current guess of the input sequence.
2. compute the Jacobians of the dynamics wrt state and input about the rolled-out sequence.

Part 1 is inherently serial: we must compute the forward pass at time t to get the state at time t + 1.

Part 2 could be either parallel or serial. We could compute it stepwise along with the forward pass. Alternatively, we could run Part 1 to completion, then compute the Jacobians wrt all (state, input) pairs along the sequence as a batch. This creates an opportunity for more SIMD speedup.

It could be faster, but we must test. It might not be worth the extra code complexity.