mlx.core.checkpoint

checkpoint(fun: Callable[P, R]) → Callable[P, R]

Transform the passed callable to one that performs gradient checkpointing with respect to the inputs of the callable.

Use this to reduce memory use for gradient computations at the expense of increased computation.

參數:

fun (Callable) -- The function to checkpoint.

回傳值:

A callable that recomputes intermediate states during gradient computation.