(lhs, rhs, dimension_numbers)
return tf.einsum(equation, lhs, rhs)
The general dot operation for TensorFlow.

  An equivalent general dot operation as that in JAX -
     <https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.dot_general.html>
  Although there is an implementation in TF XLA, avoid directly using XLA when
  possible.

  e.g., non-batched: ij,jk->ik
        batched: ijk,ikl->ijl

  Args:
    lhs: an array (the left-hand side matrix/vector to be multiplied)
    rhs: an array (the right-hand side matrix/vector to be multiplied)
    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))

  Returns:
    An array that contains the result.
The general dot operation for TensorFlow.
The general dot operation for TensorFlow .
def tf_dot_general(lhs, rhs, dimension_numbers):
  """The general dot operation for TensorFlow.

  An equivalent general dot operation as that in JAX -
     <https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.dot_general.html>
  Although there is an implementation in TF XLA, avoid directly using XLA when
  possible.

  e.g., non-batched: ij,jk->ik
        batched: ijk,ikl->ijl

  Args:
    lhs: an array (the left-hand side matrix/vector to be multiplied)
    rhs: an array (the right-hand side matrix/vector to be multiplied)
    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))

  Returns:
    An array that contains the result.
  """
  char_list = list(string.ascii_lowercase)
  char_list = char_list[8:] + char_list[:8]
  lhs_rank, rhs_rank = len(lhs.shape), len(rhs.shape)
  lhs_rep = char_list[:lhs_rank]
  rhs_rep = char_list[lhs_rank:lhs_rank + rhs_rank]
  contraction, batch = dimension_numbers
  lhs_contraction, rhs_contraction = contraction
  if len(lhs_contraction) != len(rhs_contraction):
    raise ValueError(
        "The input matrices are required to have the same number "
        "of contraction dimensions, but got: lhs {}, rhs: {}".format(
            len(lhs_contraction), len(rhs_contraction)))
  lhs_batch, rhs_batch = batch
  if len(lhs_batch) != len(rhs_batch):
    raise ValueError("The input matrices are required to have the same number "
                     "of batch dimensions, but got: lhs {}, rhs: {}".format(
                         len(lhs_batch), len(rhs_batch)))

  if not lhs_batch and not rhs_batch:
    return _non_batched_matmul(lhs, rhs, lhs_contraction, rhs_contraction)

  if (lhs_rank == rhs_rank == 3 and lhs_batch == (0,) and rhs_batch == (0,) and
      lhs_contraction == (2,) and rhs_contraction == (1,)):
    return tf.linalg.matmul(lhs, rhs)

  for i in range(len(lhs_contraction)):
    rhs_rep[rhs_contraction[i]] = lhs_rep[lhs_contraction[i]]
  for i in range(len(lhs_batch)):
    rhs_rep[rhs_batch[i]] = lhs_rep[lhs_batch[i]]

  output_rep = _compose_output_rep(lhs_rep, rhs_rep, lhs_contraction,
                                   rhs_contraction, lhs_batch, rhs_batch)
  equation = "".join(lhs_rep) + "," + "".join(rhs_rep) + "->" + output_rep
  return tf.einsum(equation, lhs, rhs)
def tf_dot_general ( lhs , rhs , dimension_numbers ) : char_list = list ( string . ascii_lowercase ) char_list = char_list [ 8 : ] + char_list [ : 8 ] lhs_rank , rhs_rank = len ( lhs . shape ) , len ( rhs . shape ) lhs_rep = char_list [ : lhs_rank ] rhs_rep = char_list [ lhs_rank : lhs_rank + rhs_rank ] contraction , batch = dimension_numbers lhs_contraction , rhs_contraction = contraction if len ( lhs_contraction ) != len ( rhs_contraction ) : raise ValueError ( "The input matrices are required to have the same number " "of contraction dimensions, but got: lhs {}, rhs: {}" . format ( len ( lhs_contraction ) , len ( rhs_contraction ) ) ) lhs_batch , rhs_batch = batch if len ( lhs_batch ) != len ( rhs_batch ) : raise ValueError ( "The input matrices are required to have the same number " "of batch dimensions, but got: lhs {}, rhs: {}" . format ( len ( lhs_batch ) , len ( rhs_batch ) ) ) if not lhs_batch and not rhs_batch : return _non_batched_matmul ( lhs , rhs , lhs_contraction , rhs_contraction ) if ( lhs_rank == rhs_rank == 3 and lhs_batch == ( 0 , ) and rhs_batch == ( 0 , ) and lhs_contraction == ( 2 , ) and rhs_contraction == ( 1 , ) ) : return tf . linalg . matmul ( lhs , rhs ) for i in range ( len ( lhs_contraction ) ) : rhs_rep [ rhs_contraction [ i ] ] = lhs_rep [ lhs_contraction [ i ] ] for i in range ( len ( lhs_batch ) ) : rhs_rep [ rhs_batch [ i ] ] = lhs_rep [ lhs_batch [ i ] ] output_rep = _compose_output_rep ( lhs_rep , rhs_rep , lhs_contraction , rhs_contraction , lhs_batch , rhs_batch ) equation = "" . join ( lhs_rep ) + "," + "" . join ( rhs_rep ) + "->" + output_rep return tf . einsum ( equation , lhs , rhs )