Pt Bmm
The batched matrix multiply behind attention, without the Python
- tens_a
- tens_b
- TENSOR
If you're here because a workflow told you to use bmm and you have no idea what it does, here's the short version: it's matrix multiplication for a whole batch of matrices at once. Where a regular matmul multiplies two 2D grids, Pt Bmm multiplies two 3D tensors of shape (batch, m, n) and (batch, n, p) and returns (batch, m, p) - the same multiply-accumulate dance, run batch times in one call.
The reason this matters in this pack specifically: batched matmul is the inner workhorse of multi-head attention. The pack has a whole "build a Transformer from scratch" guide, and attention's Q @ K^T / sqrt(d) and scores @ V steps are exactly the kind of per-batch, per-head multiplies that bmm exists for. You won't wire this into a KSampler; you'll wire it into a model-building graph.
Inputs
tens_a- a 3DTENSOR, shape(batch, m, n).tens_b- a 3DTENSOR, shape(batch, n, p), so it's compatible withtens_a.
Output: one TENSOR, shape (batch, m, p).
Where people get burned
The node is strict: both inputs must be 3D, and it says so with an explicit ValueError rather than a cryptic backend error - "torch.bmm() only supports batch matrix multiplication for 3D tensors." If your tensors are 2D, you don't want bmm at all; use the pack's PtMatMul or PtMm. And the batch dimension, the inner n dimension, and the dtypes all have to line up, or you'll get a shape mismatch deep inside torch's error messages. If a transformer graph keeps failing right around here, a PtShowSize node on both inputs is the fastest way to see what actually arrived.
Also worth knowing: there's no precision surprise here - it's standard float matmul, so use fp32 unless you've explicitly gone bf16. The pack ships PtToFloat16/PtToBfloat16 if you're chasing memory on a big model build.
Installing it
Ships in ComfyUI-Pt-Wrapper by Hide Inada (HowToSD). ComfyUI Manager → search "ComfyUI-Pt-Wrapper" → restart, or:
cd ComfyUI/custom_nodes
git clone https://github.com/HowToSD/ComfyUI-Pt-Wrapper
This pack carries a heavy requirements.txt (transformers, datasets, scikit-learn, scipy, gensim, pandas, peft, accelerate...). If Manager's install fails, pip install -r requirements.txt in the clone. No model downloads for tensor ops - only the example training workflows auto-fetch datasets like CIFAR-10.
Inputs (2)
| Name | Type | Default | Description |
|---|---|---|---|
| tens_a | TENSOR | — | |
| tens_b | TENSOR | — |
Outputs (1)
| Name | Type | Description |
|---|---|---|
| TENSOR | TENSOR | — |