code: def _dot_general(
      lhs,
      rhs,
      *,
      dimension_numbers,
      precision: tuple[PrecisionType, PrecisionType] | None,
      preferred_element_type: DType | None,
      _in_avals: Sequence[core.ShapedArray],
      _out_aval: core.ShapedArray,
  ):
      """Implementation of lax.dot_general_p in terms of tf.linalg.einsum."""
      # Unused arguments.
      del precision
      del preferred_element_type
  
      (
          lhs,
          rhs,
          convert_result,
      ) = jax2tf._dot_general_convert_to_common_dtype(
          lhs, _in_avals[0], rhs, _in_avals[1], _out_aval
      )
  
      (lhs_contracting, rhs_contracting), (
          lhs_batch,
          rhs_batch,
      ) = dimension_numbers
      lhs_ndim, rhs_ndim = len(lhs.shape), len(rhs.shape)
  
      # This condition ensures that:
      # 1) the batch dimensions are ordered in the same way in lhs and rhs (this is
      #    not strictly necessary, but we would have to reshape the array if that
      #    were not the case;
      # 2) lhs and rhs have the same number of dimensions +/- 1
      # 3) the number of non-batch dimensions in both tensors is either 1 or 2
      # 4) the contracting dimensions are consistent with those of a classic
      #    matrix/matrix, vector/matrix or matrix/vector multiplication.
      if (
          lhs_batch == rhs_batch == tuple(range(len(lhs_batch)))
          and lhs_ndim - rhs_ndim in [-1, 0, 1]
          and 1 <= lhs_ndim - len(lhs_batch) <= 2
          and 1 <= rhs_ndim - len(rhs_batch) <= 2
          and lhs_contracting == (len(lhs.shape) - 1,)
          and rhs_contracting == (len(lhs_batch),)
      ):
          # All the inputs to tf.linalg.matmul must have 2 inner dimensions,
          # after their batch dimensions, so we need to expand the dimensions
          # appropriately. We can get to this branch with three combinations of
          # inner shapes:
          # - lhs.inner_shape == [a, b], rhs.inner_shape == [b, c]
          #   - in this case, the resulting inner shape is [a, c];
          # - lhs.inner_shape == [b]   , rhs.inner_shape == [b, c]
          #   - in this case, we need to expand lhs to [1, b], and the resulting
          #     shape is [c]. We need to squeeze the result of tf.linalg.matmul
          #     as it will have shape [1, c];
          # - lhs.shape == [batch] + [a, b], rhs.shape == [batch] + [b]
          #   - in this case, we need to expand rhs to [b, 1], and the resulting
          #     shape is [a]. We need to squeeze the result of tf.linalg.matmul
          #     as it will have shape [a, 1];
          # - lhs.shape == [batch] + [b]   , rhs.shape == [batch] + [b]
          #   - in this case, we need to expand lhs to [1, b] and rhs to [b, 1],
          #     and the resulting shape is (). We need to squeeze the result of
          #     tf.linalg.matmul as it will have shape [1, 1].
          squeeze_idxs = []
          if lhs_ndim - len(lhs_batch) == 1:
              lhs = tf.expand_dims(lhs, lhs_ndim - 1)
              squeeze_idxs.append(len(lhs.shape) - 2)
          if rhs_ndim - len(rhs_batch) == 1:
              rhs = tf.expand_dims(rhs, rhs_ndim)
              squeeze_idxs.append(len(rhs.shape) - 1)
          result = tf.linalg.matmul(lhs, rhs)
          if len(squeeze_idxs) != 0:
              assert all(result.shape[i] == 1 for i in squeeze_idxs)
              result = tf.squeeze(result, squeeze_idxs)
          return convert_result(result)
  
      new_id = iter(string.ascii_letters)
      lhs_axis_ids = [next(new_id) for _ in lhs.shape]
      rhs_axis_ids = [next(new_id) for _ in rhs.shape]
      lhs_out_axis_ids = lhs_axis_ids[:]
      rhs_out_axis_ids = rhs_axis_ids[:]
  
      for lhs_axis, rhs_axis in zip(lhs_contracting, rhs_contracting):
          shared_id = next(new_id)
          lhs_axis_ids[lhs_axis] = shared_id
          rhs_axis_ids[rhs_axis] = shared_id
          lhs_out_axis_ids[lhs_axis] = None  # type: ignore[call-overload]
          rhs_out_axis_ids[rhs_axis] = None  # type: ignore[call-overload]
  
      batch_ids = []
      for lhs_axis, rhs_axis in zip(lhs_batch, rhs_batch):
          shared_id = next(new_id)
          lhs_axis_ids[lhs_axis] = shared_id
          rhs_axis_ids[rhs_axis] = shared_id
          lhs_out_axis_ids[lhs_axis] = None  # type: ignore[call-overload]
          rhs_out_axis_ids[rhs_axis] = None  # type: ignore[call-overload]
          batch_ids.append(shared_id)
  
      not_none = lambda x: x is not None
      out_axis_ids = list(
          filter(
              not_none, batch_ids + lhs_out_axis_ids + rhs_out_axis_ids
          )
      )
      assert lhs.dtype == rhs.dtype
      spec = "{},{}->{}".format(
          "".join(lhs_axis_ids),
          "".join(rhs_axis_ids),
          "".join(out_axis_ids),
      )
      return convert_result(tf.linalg.einsum(spec, lhs, rhs))
  
imports: from __future__ import annotations
  import builtins
  from collections.abc import Sequence
  import dataclasses
  from functools import partial, wraps
  import math
  import string
  from typing import Any, Callable, Optional
  from jax._src import core
  from jax import lax
  from jax._src.lax import slicing
  from jax._src import dtypes
  from jax._src import util
  from jax.experimental.jax2tf import jax2tf
  import numpy
  import tensorflow
label_desc: Use image node_id 15 for calling a global function with example usage: _dot_general(lhs, rhs) and returns: convert_result, convert_result