jax.lax.dot_general¶

jax.lax.dot_general(lhs, rhs, dimension_numbers, precision=None)[source]¶

More general contraction operator.

Wraps XLA’s DotGeneral operator.

Parameters
  • lhs (Any) – an array

  • rhs (Any) – an array

  • dimension_numbers (Tuple[Tuple[Sequence[int], Sequence[int]], Tuple[Sequence[int], Sequence[int]]]) – a tuple of tuples of the form ((lhs_contracting_dims, rhs_contracting_dims), (lhs_batch_dims, rhs_batch_dims))

  • precision (Optional[Any]) – Optional. Either None, which means the default precision for the backend, or a lax.Precision enum value (Precision.DEFAULT, Precision.HIGH or Precision.HIGHEST).

Return type

Any

Returns

An array containing the result.