CachedPartial

Create a partial from a NNX transformed function.

Abstract Signature:

CachedPartial(f: Callable)

PyTorchJAX (Core)Flax NNX

PyTorch

API: functools.partial

Strategy: Direct Mapping

Official Docs ↗

JAX (Core)

API: jax.tree_util.Partial

Strategy: Direct Mapping

Official Docs ↗

Flax NNX

API: flax.nnx.cached_partial

Strategy: Direct Mapping

Official Docs ↗