code: def dot_general(lhs, rhs, dimension_numbers):
      (lhs_contracting, rhs_contracting), (
          lhs_batch,
          rhs_batch,
      ) = dimension_numbers
      new_id = itertools.count()
      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
          rhs_out_axis_ids[rhs_axis] = None
  
      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
          rhs_out_axis_ids[rhs_axis] = None
          batch_ids.append(shared_id)
  
      not_none = lambda x: x is not None
      out_axis_ids = filter(
          not_none, batch_ids + lhs_out_axis_ids + rhs_out_axis_ids
      )
      assert lhs.dtype == rhs.dtype
      dtype = np.float32 if lhs.dtype == dtypes.bfloat16 else None
      out = np.einsum(
          lhs,
          lhs_axis_ids,
          rhs,
          rhs_axis_ids,
          out_axis_ids,
          dtype=dtype,
      )
      return (
          out.astype(dtypes.bfloat16)
          if lhs.dtype == dtypes.bfloat16
          else out
      )
  
imports: import builtins
  import collections
  import itertools
  import numpy
  import opt_einsum
  import scipy.special
  from jax._src import dtypes
  from jax._src import util
label_desc: Use image node_id 24 for calling a global function with example usage: dot_general(lhs, rhs, dimension_numbers) and returns: unknown