torch.foreach.mm

torch.foreach.mm(inputs: TensorList, mat2: TensorList, /) → tuple[Tensor, ...]

Multiplies corresponding matrices from inputs and mat2 using torch.mm(). This is semantically equivalent to applying torch.mm() independently at every list position. It does not mutate its arguments and returns a tuple of result tensors.

On supported CUDA inputs, an accelerated grouped matrix multiplication implementation may be used. Other inputs fall back to per-position execution.

Both tensor-list arguments must be non-empty and have the same length. There is no in-place torch.foreach.mm_ operation.

Parameters:

• inputs (list or tuple of Tensor) – first matrices.

• mat2 (list or tuple of Tensor) – second matrices.

Returns:

a tuple containing one matrix product per list position.