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 Reproducibility Search Size 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-08

Copy

Latest development documentation · Updated 2026-10-08

tensorplay.distributed.tensor API

Functions 11

distribute_module

functionFull reference ↗

tensorplay.distributed.tensor.distribute_module(module: Any, device_mesh: DeviceMesh | None = None, partition_fn: Callable[[str, Any, DeviceMesh], Any] | None = None, input_fn: Callable[[Any, Any, DeviceMesh], Any] | None = None, output_fn: Callable[[Any, Any, DeviceMesh], Any] | None = None) → Any[source]

distribute_tensor

functionFull reference ↗

tensorplay.distributed.tensor.distribute_tensor(tensor: Any, device_mesh: DeviceMesh | None = None, placements: Sequence[Placement] | None = None, *, src_data_rank: int | None = 0) → DTensor[source]

empty

functionFull reference ↗

tensorplay.distributed.tensor.empty(*size: Any, dtype: Any = None, layout: Any = 'strided', requires_grad: bool = False, device_mesh: DeviceMesh | None = None, placements: Sequence[Placement] | None = None) → DTensor[source]

from_local

functionFull reference ↗

tensorplay.distributed.tensor.from_local(local_tensor: Any, device_mesh: DeviceMesh | None = None, placements: Sequence[Placement] | None = None, *, run_check: bool = False, shape: Sequence[int] | None = None, stride: Sequence[int] | None = None, grad_placements: Sequence[Placement] | None = None) → DTensor

full

functionFull reference ↗

tensorplay.distributed.tensor.full(size: Any, fill_value: Any, *, dtype: Any = None, layout: Any = 'strided', requires_grad: bool = False, device_mesh: DeviceMesh | None = None, placements: Sequence[Placement] | None = None) → DTensor[source]

linspace

functionFull reference ↗

tensorplay.distributed.tensor.linspace(start: Any, end: Any, steps: int, *, dtype: Any = None, layout: Any = 'strided', requires_grad: bool = False, device_mesh: DeviceMesh | None = None, placements: Sequence[Placement] | None = None) → DTensor[source]

logspace

functionFull reference ↗

tensorplay.distributed.tensor.logspace(start: Any, end: Any, steps: int, base: float = 10.0, *, dtype: Any = None, layout: Any = 'strided', requires_grad: bool = False, device_mesh: DeviceMesh | None = None, placements: Sequence[Placement] | None = None) → DTensor[source]

ones

functionFull reference ↗

tensorplay.distributed.tensor.ones(*size: Any, dtype: Any = None, layout: Any = 'strided', requires_grad: bool = False, device_mesh: DeviceMesh | None = None, placements: Sequence[Placement] | None = None) → DTensor[source]

rand

functionFull reference ↗

tensorplay.distributed.tensor.rand(*size: Any, dtype: Any = None, layout: Any = 'strided', requires_grad: bool = False, device_mesh: DeviceMesh | None = None, placements: Sequence[Placement] | None = None) → DTensor[source]

randn

functionFull reference ↗

tensorplay.distributed.tensor.randn(*size: Any, dtype: Any = None, layout: Any = 'strided', requires_grad: bool = False, device_mesh: DeviceMesh | None = None, placements: Sequence[Placement] | None = None) → DTensor[source]

zeros

functionFull reference ↗

tensorplay.distributed.tensor.zeros(*size: Any, dtype: Any = None, layout: Any = 'strided', requires_grad: bool = False, device_mesh: DeviceMesh | None = None, placements: Sequence[Placement] | None = None) → DTensor[source]

Classes 5

DTensor

classFull reference ↗

class tensorplay.distributed.tensor.DTensor(local_tensor: Any, device_mesh: DeviceMesh, placements: Sequence[Placement], *, shape: Sequence[int] | None = None, stride: Sequence[int] | None = None, grad_placements: Sequence[Placement] | None = None, backward_dtype: Any = None): A logical tensor represented by a local value and a mesh placement.

Partial

classFull reference ↗

class tensorplay.distributed.tensor.Partial(reduce_op: str = 'sum')[source]: Store values that still need a reduction across one mesh dimension.

Placement

classFull reference ↗

class tensorplay.distributed.tensor.Placement[source]: Base class for a tensor layout on one mesh dimension.

Replicate

classFull reference ↗

class tensorplay.distributed.tensor.Replicate[source]: Keep a complete copy of the logical tensor on every rank.

Shard

classFull reference ↗

class tensorplay.distributed.tensor.Shard(dim: int)[source]: Split one logical tensor dimension across a mesh dimension.

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.distributed.rpc API

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

tensorplay.distributed.tensor.parallel API

Complete API reference for tensorplay.distributed.tensor.parallel, including signatures, parameters, examples and members.

On this page

Functions 11 distribute_module distribute_tensor empty from_local full linspace logspace ones rand randn zeros Classes 5 DTensor Partial Placement Replicate Shard