code: def _dot_general_j(
      eqn: JaxprEqn,
      idx: int,
      invals: list[ShapedArray],
      cts_in: ShapedArray,
  ) -> jnp.ndarray:
      contracting_dims, batch_dims = eqn.params["dimension_numbers"]
  
      lhs_c_dims, rhs_c_dims = contracting_dims
      lhs_b_dims, rhs_b_dims = batch_dims
  
      lhs, rhs = invals
  
      if idx == 0:
          self = lhs
          self_c_dims, self_b_dims = lhs_c_dims, lhs_b_dims
  
          other = rhs
          other_c_dims, other_b_dims = rhs_c_dims, rhs_b_dims
  
      else:
          self = rhs
          self_c_dims, self_b_dims = rhs_c_dims, rhs_b_dims
  
          other = lhs
          other_c_dims, other_b_dims = lhs_c_dims, lhs_b_dims
  
      self_ncb_dims = tuple(
          i
          for i in range(self.ndim)
          if i not in self_c_dims + self_b_dims
      )
      self_nc_dims = tuple(
          i for i in range(self.ndim) if i not in self_c_dims
      )
  
      j = jnp.moveaxis(  # pytype: disable=wrong-arg-types  # jnp-type
          other,
          other_b_dims
          + tuple(d[1] for d in sorted(zip(self_c_dims, other_c_dims))),
          tuple(range(len(other_b_dims)))
          + tuple(range(-len(other_c_dims), 0)),
      )
  
      self_ncb_out = tuple(
          utils.axis_after_dot(
              i,
              self_c_dims,
              self_b_dims,
              other.ndim if idx == 1 else None,
          )
          for i in self_ncb_dims
      )
  
      self_nc_in = tuple(cts_in.ndim + i for i in self_nc_dims)
      j = jnp.expand_dims(j, self_ncb_out + self_nc_in)
  
      self_ncb_size = utils.size_at(self, self_ncb_dims)
      self_ncb_in = tuple(i + cts_in.ndim for i in self_ncb_dims)
      shape = [1 for _ in range(j.ndim)]
      for i_out, i_in in zip(self_ncb_out, self_ncb_in):
          shape[i_out] = shape[i_in] = self.shape[i_in - cts_in.ndim]
  
      eye = jnp.eye(self_ncb_size, dtype=jnp.bool_)
      eye = eye.reshape(shape)
      j = jnp.where(eye, j, jnp.zeros((), j.dtype))
  
      for out_b, (self_b, other_b) in enumerate(
          zip(self_b_dims, other_b_dims)
      ):
          b_size = other.shape[other_b]
          eye = jnp.eye(b_size, dtype=jnp.bool_)
          shape = [1 for _ in range(j.ndim)]
          shape[out_b] = shape[cts_in.ndim + self_b] = b_size
          eye = eye.reshape(shape)
          j = jnp.where(eye, j, jnp.zeros((), j.dtype))
  
      return j
  
imports: import functools
  from typing import Any, Callable, Optional, Union
  import jax
  from jax import lax
  from jax.core import AbstractValue
  from jax.core import Jaxpr
  from jax.core import JaxprEqn
  from jax.core import Literal
  from jax.core import Primitive
  from jax.core import ShapedArray
  from jax.core import Var
  from jax.interpreters import ad
  import jax.numpy
  import numpy
  from .None import utils
  from .dataclasses import dataclass
  from .dataclasses import field
label_desc: Use image node_id 6 for calling a global function with example usage: _dot_general_j(eqn, idx, invals, cts_in) and returns: j