I meant to add that autodiff is a feature we'll be working on this upcoming spring, by using JAX arrays as the buffers. Until we have autodiff, backpropagation won't be possible.