code: def nonbatchable(
      x: PyTree,
      *,
      name: Optional[str] = None,
      msg: Optional[str] = None,
  ) -> PyTree:
      """Identity function. Raises a trace-time assert if it is batched."""
      dynamic, static = partition(x, is_array)
      flat, treedef = jtu.tree_flatten(dynamic)
      if msg is None:
          if name is None:
              name = "This operation"
          msg = f"Unexpected batch tracer. {name} cannot be vmap'd."
      bind = ft.partial(nonbatchable_p.bind, msg=msg)
      flat = map(bind, flat)
      return combine(jtu.tree_unflatten(treedef, flat), static)
  
file_path: equinox/equinox/internal/_nontraceable.py
imports: import functools
  from typing import Optional
  import jax
  import jax.core
  import jax.interpreters.ad
  import jax.interpreters.batching
  import jax.interpreters.mlir
  import jax.tree_util
  from jaxtyping import PyTree
  from .._ad import nondifferentiable
  from .._errors import error_if
  from .._filters import combine, is_array, partition
label_desc: Use image node_id 8 for calling a global function with example usage: nonbatchable(x) and returns: combine