tfep.utils.math.batchwise_dot

tfep.utils.math.batchwise_dot(x1, x2, keepdim=False)[source]

Batchwise dot product between two batches of tensors.

Parameters:

  • x1 (torch.Tensor) – A tensor of shape (*, N).

  • x2 (torch.Tensor) – A tensor of shape (*, N).

  • keepdim (bool, optional) – If True, the return value has shape (*, 1). Otherwise (*,).

Returns:

result – result[i,j,...,k] is the dot product between x1[i,j,...,k] and x2[i,j,...,k]

Return type:

torch.Tensor