Move with_memory_space_constraint primitive to jax._src.state.primitives.

PiperOrigin-RevId: 972771889

Cristian Garcia committed 1d ago

ab1e448d86a3abdb2303c97d64063b94412ea24c

Parent: 00894c5

Committed by jax authors <google-ml-automation@google.com> on 8/28/2026, 8:51:46 PM