code: def onnx_add(a, b, axis=None, broadcast=True):
      """Numpy-backed implementation of ONNX Add op."""
      if broadcast:
          axis = (a.dim - b.ndim) if axis is None else axis % a.ndim
          assert a.shape[axis:][: b.ndim] == b.shape
          b_shape = np.ones(a.ndim, dtype="int64")
          b_shape[axis : axis + b.ndim] = b.shape
          b = jnp.reshape(b, b_shape)
      return [a + b]
  
file_path: jax/examples/onnx2xla.py
imports: from io import BytesIO
  import hashlib
  import urllib.request
  import sys
  import numpy
  import onnx
  from onnx import numpy_helper
  import jax.numpy
  from jax import jit, grad
  from jax import lax
label_desc: Use image node_id 4 for calling a global function with example usage: onnx_add(a, b, axis, broadcast) and returns: list