Skip to contents

Primitive Dot General

Source: R/primitives.R
prim_dot_general.Rd

General dot product of two arrays, supporting contraction over arbitrary axes and batching.

Usage

prim_dot_general(
  lhs,
  rhs,
  contracting_axes,
  batching_axes,
  precision = "highest"
)

Arguments

lhs, rhs

(arrayish)
Two inputs of the same data type whose shapes are constrained by contracting_axes and batching_axes rather than having to match. Can be any data type. R values assume the other operand's data type when it is in their data type category, and their default data type when neither operand has one.

contracting_axes

(list(integer(), integer()))
A list of two integer vectors specifying which axes of lhs and rhs to contract over. The contracted axes must have matching sizes.

batching_axes

(list(integer(), integer()))
A list of two integer vectors specifying which axes of lhs and rhs are batch axes. These must have matching sizes.

precision

(character(1))
Controls the trade-off between speed and numerical accuracy of the operation. One of "highest" (default), "high" or "default". Only the StableHLO backend honors this; it is ignored by the quickr backend.

Value

(arrayish)
Has the data type the operands agreed on. The output shape is the batch axes followed by the remaining (non-contracted, non-batched) axes of lhs, then rhs.

Implemented Rules

  • stablehlo

  • quickr

  • reverse

StableHLO

Lowers to hlo_dot_general(), specified under dot_general.

See also

nv_matmul(), %*%

Examples

# contracting a 2x3 with a 3x2 gives a 2x2 at the operands' data type
x <- nv_matrix(1:6, nrow = 2)
y <- nv_matrix(1:6, nrow = 3)
prim_dot_general(x, y,
  contracting_axes = list(2L, 1L),
  batching_axes = list(integer(0), integer(0))
)
#> AnvlArray
#>  22 49
#>  28 64
#> [ CPUi32{2,2} ]