ComfyUI Node

Pt Einsum

Pt Einsum — the one node that replaces half the tensor-math drawer

By HowToSD·Created about a year ago·Updated about a year ago· 7
Pt Einsum
  • tens_a
  • tens_b
  • tens_c
  • TENSOR
equation

Einsum is Einstein summation notation - a compact way to say "multiply these tensors and sum over these axes," and it covers matrix multiply, dot products, transposes, traces, batch matmuls, and a dozen other operations with one syntax. PtEinsum brings it to ComfyUI: type the equation, feed the tensors, get the result. If you know even a little einsum, this one node replaces half the pack's arithmetic drawer. If you don't, it's still worth learning here because the payoff is a single node instead of five.

It's part of ComfyUI-Pt-Wrapper, HowToSD's no-code PyTorch toolkit. It's the node the pack's own docs point you to for "matrix operations" - the natural choice when you're building linear algebra into a model pipeline without writing code.

How it works

It's a straight pass-through to torch.einsum. You provide an equation string and one to three tensors; the node routes based on how many you connected. One tensor: torch.einsum(eq, tens_a). Two: torch.einsum(eq, tens_a, tens_b). Three: all three. The author's one real constraint is in the docstring: sublist format is not supported - so the fancy (b,h,s,d)->... bracket form won't work; stick to the classic letter-notation like "ij,jk->ik".

Inputs that matter

  • equation - the Einstein notation string. The classic matmul is "ij,jk->ik" (tens_a is i×j, tens_b is j×k, result is i×k). Trace: "ii->". Transpose: "ij->ji". Batch matmul: "bij,bjk->bik".
  • tens_a - required first tensor.
  • tens_b, tens_c - optional, wire only if your equation uses them.

Output is a single TENSOR whose shape is determined by your equation's right-hand side.

Installing it

ComfyUI Manager → Install Custom Nodes → search "Pt Wrapper" → install → restart, or:

cd ComfyUI/custom_nodes
git clone https://github.com/HowToSD/ComfyUI-Pt-Wrapper

Restart ComfyUI after. First install is heavy (matplotlib, pandas, scipy, scikit-learn, transformers, datasets, pinned gensim), but this node is pure torch.

Common issues

The equation string is the whole game, and it's also where everything breaks. A typo - wrong letter, missing arrow, unmatched dimension name - throws an einsum error the moment you run it, and the message can be cryptic. The most common beginner mistake is swapping operands: einsum is order-sensitive, so "ij,jk->ik" is matmul but "ij,kj->ik" is something else entirely. Start with the three classics above, verify the output shape with the pack's PtShowSize, and only then get fancy. And remember: no sublist brackets.

CategoryData Analysis

Inputs (4)

NameTypeDefaultDescription
equationSTRING
tens_aTENSOR
tens_boptTENSOR
tens_coptTENSOR

Outputs (1)

NameTypeDescription
TENSORTENSOR