jax

affe2e73 - Rename `dot_with_no_batch_dims_saveable` to `dots_with_no_batch_dims_saveable` for internal consistency