code: def _reshape_axis_into(src, dst, x):
      # NB: `dst` is the number of the dimension that we should reshape into
      # *after* `src` is removed from `x`'s list of dimensions. For example, if
      # `src` is an added batch dimension, `dst` might name a target dimension in
      # the unbatched list of dimensions.
      perm = [i for i in range(x.ndim) if i != src]
      perm.insert(dst, src)
      new_shape = list(np.delete(x.shape, src))
      new_shape[dst] *= x.shape[src]
      return lax.reshape(x, new_shape, perm)
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  from functools import partial
  import operator
  from typing import NamedTuple, Union
  import numpy
  from jax._src import core
  from jax._src import dtypes
  from jax._src import util
  from jax._src.interpreters import ad
  from jax._src.interpreters import batching
  from jax._src.interpreters import mlir
  from jax._src.lax import lax
  from jax._src.lib.mlir.dialects import hlo
  from jax._src.typing import Array, DTypeLike
label_desc: Use image node_id 14 for calling a global function with example usage: _reshape_axis_into(src, dst, x) and returns: lax