def METHOD_NAME(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,
  )