jax.lax.dot_general#

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

More general contraction operator.

Wraps XLA’s DotGeneral operator.

Parameters
Return type

Array

Returns

An array containing the result.