jax.checkpoint_policies.dots_with_no_batch_dims_saveable

checkpoint_policies.dots_with_no_batch_dims_saveable = DotsSaveable(only_if_no_batch_dims=True)