code: def test_jax_to_hlo_axpy(self):
      hlo_proto, hlo_text = jax_to_ir.jax_to_hlo(
          axpy,
          [
              ("y", jax_to_ir.parse_shape_str("f32[128]")),
              ("a", jax_to_ir.parse_shape_str("f32[]")),
              ("x", jax_to_ir.parse_shape_str("f32[128,2]")),
          ],
      )
  
      # Check that hlo_text contains a broadcast, add, and multiply.
      self.assertIn("broadcast", hlo_text)
      self.assertIn("add", hlo_text)
      self.assertIn("multiply", hlo_text)
  
      # Check that the HLO parameters are in the order we specified in the
      # jax_to_hlo call.
      self.assertIn("f32[128]{0} parameter(0)", hlo_text)
      self.assertIn("f32[] parameter(1)", hlo_text)
      self.assertIn("f32[128,2]{1,0} parameter(2)", hlo_text)
  
      # Check that the parameters are in the expected order.
  
      # TODO(jlebar): Ideally we'd check that hlo_proto can be deserialized to a
      # valid HLO proto, but we don't seem to have access to hlo_pb2 at the
      # moment, so the best we seem to be able to do is check that it's nonempty.
      assert hlo_proto
  
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 the JaxToIRTest obj's underlying member method code with example usage: obj.test_jax_to_hlo_axpy() without return types