Website Blog TensorPlay documentation

tensorplay API tensorplay._composite_funcs API tensorplay._einsum API tensorplay._shape_funcs API tensorplay._stax.api API tensorplay.accelerator API tensorplay.amp.autocast_mode API tensorplay.ao.nn.intrinsic API tensorplay.ao.nn.qat API tensorplay.ao.nn.quantized API tensorplay.ao.nn.quantized.dynamic API tensorplay.ao.pruning API tensorplay.ao.quantization API tensorplay.autograd API tensorplay.autograd.anomaly_mode API tensorplay.autograd.forward_ad API tensorplay.autograd.function API tensorplay.autograd.functional API tensorplay.autograd.grad_mode API tensorplay.autograd.gradcheck API tensorplay.autograd.graph API tensorplay.autograd.profiler_util API tensorplay.autograd.variable API tensorplay.backends.cpu API tensorplay.backends.cuda API tensorplay.backends.mkl API tensorplay.backends.nnpack API tensorplay.backends.openmp API tensorplay.compiler API tensorplay.cuda API tensorplay.cuda.gds API tensorplay.cuda.graph_annotations API tensorplay.cuda.graphs API tensorplay.cuda.green_contexts API tensorplay.cuda.memory API tensorplay.cuda.nccl API tensorplay.cuda.nvtx API tensorplay.cuda.profiler API tensorplay.cuda.random API tensorplay.cuda.streams API tensorplay.cuda.tunable API tensorplay.distributed.algorithms API tensorplay.distributed.algorithms.ddp_comm_hooks.ddp_zero_hook API tensorplay.distributed.algorithms.ddp_comm_hooks.default_hooks API tensorplay.distributed.algorithms.ddp_comm_hooks.post_localSGD_hook API tensorplay.distributed.algorithms.ddp_comm_hooks.powerSGD_hook API tensorplay.distributed.algorithms.ddp_comm_hooks.quantization_hooks API tensorplay.distributed.algorithms.model_averaging API tensorplay.distributed.autograd API tensorplay.distributed.checkpoint API tensorplay.distributed.device_mesh API tensorplay.distributed.elastic.agent.server API tensorplay.distributed.elastic.events API tensorplay.distributed.elastic.metrics API tensorplay.distributed.elastic.multiprocessing API tensorplay.distributed.elastic.rendezvous API tensorplay.distributed.elastic.timer API tensorplay.distributed.elastic.utils API tensorplay.distributed.elastic.utils.data API tensorplay.distributed.fsdp API tensorplay.distributed.fsdp.sharded_grad_scaler API tensorplay.distributed.launcher API tensorplay.distributed.nn API tensorplay.distributed.nn.api API tensorplay.distributed.nn.functional API tensorplay.distributed.optim API tensorplay.distributed.pipelining API tensorplay.distributed.rpc API tensorplay.distributed.tensor API tensorplay.distributed.tensor.parallel API tensorplay.distributions API tensorplay.export API tensorplay.fft._frequency API tensorplay.fft._transforms API tensorplay.func API tensorplay.functional API tensorplay.futures API tensorplay.library API tensorplay.linalg API tensorplay.masked API tensorplay.monitor API tensorplay.multiprocessing API tensorplay.nested API tensorplay.nn.attention API tensorplay.nn.attention.bias API tensorplay.nn.attention.varlen API tensorplay.nn.functional API tensorplay.nn.init API tensorplay.nn.modules.flatten API tensorplay.nn.modules.lazy API tensorplay.nn.modules.module API tensorplay.nn.parameter API tensorplay.nn.utils.rnn API tensorplay.nn.utils.stateless API tensorplay.onnx API tensorplay.overrides API tensorplay.package API tensorplay.profiler API tensorplay.profiler.profiler API tensorplay.random API tensorplay.serialization API tensorplay.serialization.api API tensorplay.signal.windows API tensorplay.special._elementwise API tensorplay.special._error API tensorplay.special._gamma API tensorplay.stax API tensorplay.testing API tensorplay.utils.checkpoint API tensorplay.utils.data API tensorplay.utils.viz API

Automatic differentiation package - tensorplay.autograd Automatic Mixed Precision package - tensorplay.amp Complex numbers Computation-graph visualization - tensorplay.utils.viz DDP Communication Hooks Distributed communication package - tensorplay.distributed DLPack Environment variables foreach operations Index Meta device Multiprocessing package - tensorplay.multiprocessing Patching batch norm Pruning - tensorplay.ao.pruning Python Module Index Quantization - tensorplay.ao.quantization Quantum simulation Reproducibility Search Size Spiking neural networks Storage Tensor attributes Tensor views tensorplay tensorplay._stax tensorplay.accelerator tensorplay.audio tensorplay.backends tensorplay.compiler tensorplay.cuda tensorplay.cuda.tunable tensorplay.distributed.algorithms tensorplay.distributed.autograd tensorplay.distributed.checkpoint tensorplay.distributed.device_mesh tensorplay.distributed.elastic tensorplay.distributed.fsdp tensorplay.distributed.fsdp.fully_shard tensorplay.distributed.nn tensorplay.distributed.optim tensorplay.distributed.pipelining tensorplay.distributed.rpc tensorplay.distributed.tensor tensorplay.distributed.tensor.parallel tensorplay.distributions tensorplay.export tensorplay.fft tensorplay.func tensorplay.func API reference tensorplay.func whirlwind tour tensorplay.futures tensorplay.hub tensorplay.library tensorplay.linalg tensorplay.masked tensorplay.monitor tensorplay.nested tensorplay.nn tensorplay.nn.attention tensorplay.nn.attention.bias tensorplay.nn.attention.omni_attention tensorplay.nn.attention.varlen tensorplay.nn.functional tensorplay.nn.init tensorplay.onnx tensorplay.optim tensorplay.overrides tensorplay.package tensorplay.profiler tensorplay.random tensorplay.serialization tensorplay.signal tensorplay.sparse tensorplay.special tensorplay.stax tensorplay.testing tensorplay.utils.checkpoint tensorplay.utils.data tensorplay.vision Type information UX limitations

API reference

Documentation/API reference

latest (dev)

latest (dev)2026-10-15

Copy

Latest development documentation · Updated 2026-10-15

tensorplay.utils.checkpoint API

Functions 6

checkpoint_sequential

functionFull reference ↗

tensorplay.utils.checkpoint.checkpoint_sequential(function: Callable[[...], Any], chunks: int, *args: Any, **kwargs: Any) → Any[source]

checkpoint

functionFull reference ↗

tensorplay.utils.checkpoint.checkpoint(function: Callable[[...], Any] | None = None, *args: Any, use_reentrant: bool | None = None, preserve_rng_state: bool = True, context_fn: Callable[[], Any] | None = None, determinism_check: str = 'default', debug: bool = False, early_stop: bool | None = None, respect_saved_tensors_hooks: bool | None = None, **kwargs: Any) → Any[source]

create_selective_checkpoint_contexts

functionFull reference ↗

tensorplay.utils.checkpoint.create_selective_checkpoint_contexts(policy_fn_or_list: Callable[[...], Any] | list[Any], allow_cache_entry_mutation: bool = False)[source]

set_checkpoint_debug_enabled

functionFull reference ↗

tensorplay.utils.checkpoint.set_checkpoint_debug_enabled(enabled: bool | None)[source]

set_checkpoint_early_stop

functionFull reference ↗

tensorplay.utils.checkpoint.set_checkpoint_early_stop(enabled: bool)[source]

set_device_states

functionFull reference ↗

tensorplay.utils.checkpoint.set_device_states(devices: Iterable[int], states: Iterable[Any], *, device_type: str | None = None) → None[source]

Classes 3

CheckpointPolicy

classFull reference ↗

class tensorplay.utils.checkpoint.CheckpointPolicy(*values)[source]

GraphExecGroup

classFull reference ↗

class tensorplay.utils.checkpoint.GraphExecGroup[source]

SelectiveCheckpointContext

classFull reference ↗

class tensorplay.utils.checkpoint.SelectiveCheckpointContext(*, is_recompute: bool, op_output: Any = None)[source]

Help improve this page

Found an error, an unclear step, or a missing example?

Was this page helpful?

Suggest an edit Report an issue

tensorplay.testing API

Complete API reference for tensorplay.testing, including signatures, parameters, examples and members.

tensorplay.utils.data API

Complete API reference for tensorplay.utils.data, including signatures, parameters, examples and members.

On this page

Functions 6 checkpoint_sequential checkpoint create_selective_checkpoint_contexts set_checkpoint_debug_enabled set_checkpoint_early_stop set_device_states Classes 3 CheckpointPolicy GraphExecGroup SelectiveCheckpointContext