jax

db9bbb35 - Add `_dot_general_batch_rule` back because some internal projects depend on it.