code: def _check_output_dims(
      func: Callable,
      dim_sizes: dict[str, int],
      expected_output_core_dims: list[CoreDims],
      error_context: str = "",
  ) -> Callable:
      """Check that output core dimensions match the signature."""
  
      def wrapped(*args):
          out = func(*args)
          out_shapes = map(
              jnp.shape, out if isinstance(out, tuple) else [out]
          )
  
          if expected_output_core_dims is None:
              output_core_dims = [()] * len(out_shapes)
          else:
              output_core_dims = expected_output_core_dims
              if len(output_core_dims) > 1 and not isinstance(
                  out, tuple
              ):
                  raise TypeError(
                      "output must be a tuple when multiple outputs are expected, "
                      "got: {!r}\n{}".format(out, error_context)
                  )
              if len(out_shapes) != len(output_core_dims):
                  raise TypeError(
                      "wrong number of output arguments: expected %r, got %r %s"
                      % (
                          len(output_core_dims),
                          len(out_shapes),
                          error_context,
                      )
                  )
  
          sizes = dict(dim_sizes)
          for shape, core_dims in zip(out_shapes, output_core_dims):
              _update_dim_sizes(
                  sizes, shape, core_dims, error_context, is_input=False
              )
  
          return out
  
      return wrapped
  
file_path: jax/jax/_src/numpy/vectorize.py
imports: from __future__ import annotations
  from collections.abc import Collection, Sequence
  import functools
  import re
  from typing import Any, Callable
  from jax._src import api
  from jax import lax
  from jax._src.numpy import lax_numpy
  from jax._src.util import safe_map, safe_zip
label_desc: Use image node_id 4 for calling a global function with example usage: _check_output_dims(func, dim_sizes, expected_output_core_dims, error_context) and returns: wrapped, out