DistributedBinaryContraction#

class cuquantum.tensornet.experimental.distributed.DistributedBinaryContraction( expr, a: DistributedTensor, b: DistributedTensor, *, c: DistributedTensor | None = None, out: BlockCyclic, qualifiers=None, stream=None, options=None, execution=None, )[source]#

Stateful distributed binary contraction :math:D = \alpha A B + \beta C.

Construct the problem, call plan(), then call execute() one or more times. Release resources with free() or a context manager.

Inputs are DistributedTensor objects. out is a BlockCyclic distribution describing the output placement; output storage is allocated freshly by each execute() call.

Currently, this stateful form requires compact Fortran-order local input shards and does not stage incompatible layouts. Allocate shards with empty_local() or use asfortranarray. The single-use distributed_binary_contraction() stages incompatible inputs automatically.

Operand layouts are fixed at construction. Replacements supplied through reset_operands() must have matching global and local shapes, dtype, element strides, and distribution.

Warning

This API is experimental and subject to future changes.

Note

plan() performs Python-side validation before entering collective operations. Pass consistent arguments on every rank, or set options.collective_error_agreement=True to make these pre-checks fail consistently across ranks.

Parameters:

expr – Explicit binary einsum expression with two inputs and one output.
a – First distributed input tensor.
b – Second distributed input tensor.
c – Optional distributed addend with the output’s global shape.
out – Distribution describing the output placement.
qualifiers – Reserved; currently unsupported.
stream – Accepted for interface consistency; construction performs no device work. Pass a stream to plan() and execute().
options – A DistributedContractionOptions object or dict; options.handle is required by plan().
execution – Reserved; currently unsupported.

Methods

__init__( expr, a: DistributedTensor, b: DistributedTensor, *, c: DistributedTensor | None = None, out: BlockCyclic, qualifiers=None, stream=None, options=None, execution=None, )[source]#

execute( *, alpha=1.0, beta=None, release_workspace=False, stream=None, )[source]#

Execute the planned contraction.

Parameters:

alpha – Scale factor for the A @ B product.
beta – Scale factor for the addend c. Must be set when c was specified at construction, and must be left unset (the default) otherwise.
release_workspace – If True, release the workspace memory back to the allocator after this call; otherwise it is kept for a subsequent execute() call.
stream – CUDA stream used for output allocation and execution.

Returns:

The freshly allocated output DistributedTensor.

free()[source]#: Release all resources associated with this operation.

plan(*, stream=None)[source]#

Plan the contraction and allocate workspace storage.

When options.collective_error_agreement=True, rank-local pre-checks are agreed across the initialized distributed process group before collective planning begins.

Parameters:: stream – CUDA stream used for workspace setup.

release_operands()[source]#: Release operand and output references while retaining any prepared plan.

reset_operands( *, a: DistributedTensor | None = None, b: DistributedTensor | None = None, c: DistributedTensor | None = None, )[source]#

Attach replacement operands without changing the problem layout.

Each replacement must match its construction-time global and local shapes, dtype, element strides, and distribution. c may be replaced only when the problem was constructed with an addend.

After release_operands(), all originally provided operands must be supplied together. The output placement remains fixed.

Parameters:

a – Replacement first input tensor.
b – Replacement second input tensor.
c – Replacement addend tensor.