code: def axis_after_dot(
      axis: int,
      contracting_dims: Sequence[int],
      batch_dims: Sequence[int],
      lhs_ndim: Optional[int] = None,
  ) -> int:
      if axis in batch_dims:
          return batch_dims.index(axis)
  
      return (
          axis
          - sum(1 for i in contracting_dims if i < axis)
          + sum(1 for i in batch_dims if i > axis)
          + (
              0
              if lhs_ndim is None
              else lhs_ndim - len(batch_dims) - len(contracting_dims)
          )
      )
  
imports: from collections import namedtuple
  import functools
  import inspect
  import operator
  import types
  from typing import Any, Callable, Iterable, Optional, Sequence, Sized, TypeVar, Union
  import warnings
  import jax
  from jax import core
  from jax import random
  import jax.numpy
  from jax.tree_util import tree_all
  from jax.tree_util import tree_map
  import numpy
label_desc: Use image node_id 24 for calling a global function with example usage: axis_after_dot(axis, contracting_dims, batch_dims, lhs_ndim) and returns: unknown, batch_dims