code: def concise_vmap(
      module, in_axes, out_axes, axis_size=None, **var_specs
  ):
      variable_axes = {
          k: v[0]
          for k, v in var_specs.items()
          if isinstance(v, Sequence)
      }
      splits = {
          k: v[1]
          for k, v in var_specs.items()
          if isinstance(v, Sequence)
      }
      return vmap(
          module,
          in_axes=in_axes,
          out_axes=out_axes,
          variable_axes=variable_axes,
          split_rngs=splits,
          axis_size=axis_size,
      )
  
imports: import functools
  from pprint import pprint
  from typing import Any, Callable, Optional, Sequence
  from flax.core.frozen_dict import unfreeze
  from flax.linen import initializers
  from flax.linen import Module, compact, vmap
  from flax.linen.linear import PrecisionLike
  import jax
  from jax import lax, numpy, random
label_desc: Use image node_id 1 for calling a global function with example usage: concise_vmap(module, in_axes, out_axes, axis_size) and returns: vmap