Torch/JAX don't share many parallels when it comes to padding in Haiku - reconfiguring how to implement padding for JAX resents