code: def axpy(a, x, y):
      return a * x + y[:, jnp.newaxis]
  
file_path: jax/tests/jax_to_ir_test.py
imports: import unittest
  from absl.testing import absltest
  import jax.numpy
  from jax.tools import jax_to_ir
  from jax._src import test_util
label_desc: Use image node_id 1 for calling a global function with example usage: axpy(a, x, y) and returns: unknown