Tensor Contractions

For Einstein-index notation, supported operations, and scalar handling, see the TensorOperations.jl guide.

Low-level contract / contract!

Prefer `@tensor`. `contract` and `contract!` are the pairwise primitive the extension calls. They do not parse einsum strings, pick a contraction order, or free intermediates for you.

Mode labels are not einsum strings: "ik" with "kj" is a GEMM. "ijk" with "ikl" and explicit output "ijl" is a batched product. Labels may be ASCII strings, Char tuples, or integers (1 maps to 'a').

A = cuNumeric.rand(Float32, 64, 32)
B = cuNumeric.rand(Float32, 32, 16)
C = contract(A, "ik", B, "kj")              # allocates, Einstein output order
contract!(similar(C), "ij", A, "ik", B, "kj"; α=2, β=0)

# batched: keep the shared 'i' on the output
AA = cuNumeric.rand(Float32, 8, 16, 32)
BB = cuNumeric.rand(Float32, 8, 32, 4)
CC = cuNumeric.zeros(Float32, 8, 16, 4)
contract!(CC, "ijl", AA, "ijk", BB, "ikl")

tensordot(AA, BB, ([3], [2]))               # same contraction, axes form

α and β implement C = β*C + α*(A ⋆ B) in Julia. Prefer a Julia Number. A 0D NDArray is also accepted. The C++ kernel always writes the unscaled product. β ≠ 0 uses a temporary. Duplicate labels inside one array are not allowed. Use cuNumeric.diagonal first.

This path is multi-GPU via Legate tiling (per-tile cuTENSOR or TBLIS), not cuTensorMp. Integer and Bool inputs promote to Float64 like other linalg.

cuNumeric.contract! Method

contract!(C, Cmodes, A, Amodes, B, Bmodes; α=1, β=0)

In-place pairwise tensor contraction C = β * C + α * (A ⋆ B).

α and β may be a Julia Number or a 0-d NDArray.

Amodes, Bmodes, and Cmodes are mode labels for A, B, and C: an ASCII AbstractString, a tuple/vector of Char, or a vector of integers (1 maps to 'a'). Each label appears twice (a contracted or free index) or three times (a batched / Hadamard index). Duplicate labels inside one array are not allowed — extract a diagonal first with cuNumeric.diagonal.

Supported element types are Float32, Float64, ComplexF32, and ComplexF64. Integer and Bool inputs are converted to Float64 under the usual promotion rules.

This is the pairwise primitive TensorOperations.jl can call later. It does not parse einsum strings. Multi-GPU execution uses Legate tiling (not cuTensorMp).

source

cuNumeric.contract Method

contract(A, Amodes, B, Bmodes; α=1)

Allocate and return α * (A ⋆ B). Output modes are the labels that appear once, in the order they occur on A then B (classical Einstein). For a batched or Hadamard product, allocate C yourself and call contract! with explicit Cmodes.

source

cuNumeric.tensordot Function

tensordot(A, B, axes=2; α=1)
tensordot(A, B, (a_axes, b_axes); α=1)

Contract A with B along the given 1-based axes. axes::Integer contracts the last axes dimensions of A with the first axes of B. A tuple of axis collections names the axes on each input. Remaining axes of A then B become the output.

source