UpdateContext ============= Context manager/decorator for handling complex state updates with graph mutations. **Abstract Signature:** ``UpdateContext(tag: str = 'default')`` .. raw:: html

PyTorch

API:
Strategy: Custom / Partial

JAX (Core)

API:
Strategy: Custom / Partial

Flax NNX

API: flax.nnx.update_context
Strategy: Direct Mapping