code: def _argnum_weak_type(*argnums):
      return lambda *args, **_: all(args[i].weak_type for i in argnums)
  
file_path: jax/jax/_src/lax/utils.py
imports: from functools import partial
  import operator
  from jax._src import core
  from jax._src import dispatch
  from jax._src import dtypes
  from jax._src.util import safe_zip
  from jax._src.lib import xla_client
  import numpy
label_desc: Use image node_id 2 for calling a global function with example usage: _argnum_weak_type() and returns: function