fn

addbmm

Tensor
addbmm(input: Tensor, batch1: Tensor, batch2: Tensor, beta: float = ..., alpha: float = ...)
source

Batched matmul with reduction over the batch axis.

Computes βinput+αk(batch1[k]batch2[k])\beta \cdot \text{input} + \alpha \cdot \sum_{k} (\text{batch1}[k] \cdot \text{batch2}[k]). Useful for accumulating multiple parallel matrix products into a single output matrix (e.g. summing per-head attention contributions).

Parameters

inputTensor
Accumulator of shape (M, N).
batch1Tensor
Batched left matrices of shape (B, M, K).
batch2Tensor
Batched right matrices of shape (B, K, N).
betafloat
Scalar multiplier on input. Defaults to 1.0.
alphafloat
Scalar multiplier on the batched matmul sum. Defaults to 1.0.

Returns

Tensor

Tensor of shape (M, N).

Notes

Mathematical definition:

out=βinput+αk=0B1(batch1[k]batch2[k]).\text{out} = \beta \cdot \text{input} + \alpha \sum_{k = 0}^{B - 1} (\text{batch1}[k] \cdot \text{batch2}[k]).

Differs from baddbmm in that the batch dimension is reduced away; baddbmm keeps the per-batch results.

Examples

>>> import lucid
>>> b1 = lucid.ones((3, 2, 4))
>>> b2 = lucid.ones((3, 4, 2))
>>> M = lucid.zeros((2, 2))
>>> lucid.addbmm(M, b1, b2)
Tensor([[12., 12.],
        [12., 12.]])