Hi,
Recently, I met the same problem and tried to do the batch_jacobian operation with for-loops. Although it works, but the runtime is too long. Fianlly,  I implemented the batch_jacobian in another way, it is more efficient and the runtime is close to the tf.GradientTape.
def batch_jacobian(func, x, create_graph=False):
  # x in shape (Batch, Length)
  def _func_sum(x):
    return func(x).sum(dim=0)
  return autograd.functional.jacobian(_func_sum, x, create_graph=create_graph).permute(1,0,2)