code: def common_rewrite(
      cond_fun,
      body_fun,
      init_val,
      max_steps,
      buffers,
      makes_false_steps,
  ):
      """Handles:
  
      - Efficient in-place updates;
      - Efficient/correct vmap;
      - max_steps.
  
      The efficient in-place updates are done using buffers. These mean that we get to
      interchange `select` (from vmap) and `scatter` (from the in-place update) so that
      they happen in the efficient order. Specifically for `checkpointed_while_loop`, it
      also means we don't need to save our buffers in the list of checkpoints.
  
      The efficient vmap is done by always having an unvmap'd return value from
      `cond_fun`: essentially we reduce from the vmap'd case to the no-vmap case here.
      This ensures that (a) `bounded_while_loop` exhibits early-exit, as our `lax.cond`s
      don't get turned into `lax.select`s, and (b) `checkpointed_while_loop` has the
      correct logic on the backward pass. (As all checkpoints need to be in lockstep
      across the batch.)
      """
  
      if buffers is None:
          new_buffers = lambda is_leaf: lambda val: []
      else:
  
          def new_buffers(is_leaf):
              def new_buffers2(val):
                  _, _, _, val = val  # ignore step and pred
                  bufs = buffers(val)
                  # Ignore the ._pred attribute of nested buffers.
                  # This is kind of a hack: we're special-casing support for nested
                  # buffers so that end users don't have to do the ._array lookup
                  # themselves.
                  tree = jtu.tree_map(
                      _unwrap_buffers, bufs, is_leaf=_is_buffer
                  )
                  return jtu.tree_leaves(tree, is_leaf=is_leaf)
  
              return new_buffers2
  
      def _wrap_buffers(val, pred, tag):
          def wrap_buffer(leaf):
              if not is_array(leaf):
                  raise ValueError(
                      "Only arrays can be treated as buffers."
                  )
              return _Buffer(leaf, pred, tag, makes_false_steps)
  
          _, _, _, buffer_val = tree_at(
              new_buffers(None),
              (None, None, None, val),
              replace_fn=wrap_buffer,
          )
          return buffer_val
  
      def new_cond_fun(val):
          step, pred, prev_pred, val = val
          del pred
          # Note that this is actually recomputing `pred`!
          # For some reason this is actually a minor performance optimisation.
          # See https://github.com/patrick-kidger/diffrax/issues/274
          buffer_val = _wrap_buffers(val, prev_pred, None)
          out = unvmap_any(cond_fun(buffer_val))
          if max_steps is not None:
              if type(max_steps) is not int:
                  raise ValueError(
                      "`max_steps` must be a Python integer"
                  )
              out = out & (step < max_steps)
          return nonbatchable(out)
  
      def new_body_fun(val):
          tag = object()
  
          def is_our_buffer(node):
              return isinstance(node, _Buffer) and node._tag is tag
  
          def unwrap_and_select(leaf, leaf2):
              if is_our_buffer(leaf):
                  assert is_our_buffer(leaf2)
                  # sanity check that this is the lowest buffer, i.e. when nesting
                  # multiple checkpointed_while_loops.
                  assert is_array(leaf._array)
                  assert is_array(leaf2._array)
                  return leaf2._array
              else:
                  return _select_if_vmap(
                      pred, leaf2, leaf, makes_false_steps
                  )
  
          step, pred, _, val = val
          buffer_val = _wrap_buffers(val, pred, tag)
          buffer_val2 = body_fun(buffer_val)
          # Needed to work with `disable_jit`, as then we lose the automatic
          # ArrayLike->Array cast provided by JAX's while loops.
          # The input `val` is already cast to Array below, so this matches that.
          buffer_val2 = jtu.tree_map(fixed_asarray, buffer_val2)
          # Strip `.named_shape`; c.f. Diffrax issue #246
          struct = jax.eval_shape(lambda: buffer_val)
          struct2 = jax.eval_shape(lambda: buffer_val2)
          struct = jtu.tree_map(
              lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype), struct
          )
          struct2 = jtu.tree_map(
              lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype), struct2
          )
          if not tree_equal(struct, struct2):
              string = tree_pformat(struct, struct_as_array=True)
              string2 = tree_pformat(struct2, struct_as_array=True)
              out = []
              for line, line2 in it.zip_longest(
                  string.split("\n"), string2.split("\n"), fillvalue=""
              ):
                  if line == line2:
                      out.append("  " + line)
                  else:
                      out.append("- " + line)
                      out.append("+ " + line2)
              out = "\n".join(out)
              raise ValueError(
                  "`body_fun` must have the same input and output structure. Difference "
                  "is:\n" + out
              )
          val2 = jtu.tree_map(
              unwrap_and_select,
              buffer_val,
              buffer_val2,
              is_leaf=is_our_buffer,
          )
          step2 = step + 1
          pred2 = pred & cond_fun(buffer_val2)
          return step2, pred2, pred, val2
  
      init_val = jtu.tree_map(fixed_asarray, init_val)
      init_buffer_val = _wrap_buffers(init_val, jnp.array(True), None)
      init_pred = jnp.array(cond_fun(init_buffer_val))
      new_init_val = (
          jnp.array(0),
          init_pred,
          jnp.array(True),
          init_val,
      )
  
      return new_cond_fun, new_body_fun, new_init_val, new_buffers
  
imports: import itertools
  from typing import Any, TYPE_CHECKING, Union
  import jax
  import jax.core
  import jax.interpreters.ad
  import jax.interpreters.batching
  import jax.interpreters.mlir
  import jax.lax
  import jax.numpy
  import jax.tree_util
  from jaxtyping import Array, Bool
  from ..._filters import combine, is_array, partition
  from ..._module import field, Module
  from ..._pretty_print import tree_pformat
  from ..._tree import tree_at, tree_equal
  from ..._unvmap import unvmap_any
  from .._nontraceable import nonbatchable
  from .._primitive import create_vprim
label_desc: Use image node_id 15 for calling a global function with example usage: common_rewrite(cond_fun, body_fun, init_val, max_steps, buffers, makes_false_steps) and returns: new_cond_fun,  new_body_fun,  new_init_val,  new_buffers, buffer_val, nonbatchable, step2,  pred2,  pred,  val2, new_buffers2, _Buffer, unknown, jtu, leaf2, _select_if_vmap