def METHOD_NAME(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]