jax.checkpoint_policies.checkpoint_dots_with_no_batch_dims

checkpoint_policies.checkpoint_dots_with_no_batch_dims = DotsSaveable(only_if_no_batch_dims=True)