code: def _reshape_axis_out_of(src, size1, x):
      shape = list(x.shape)
      size2, ragged = divmod(shape[src], size1)
      assert not ragged
      shape[src : src + 1] = [size1, size2]
      return lax.reshape(x, shape)
  
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 15 for calling a global function with example usage: _reshape_axis_out_of(src, size1, x) and returns: lax