code: def test_jax_to_tf_axpy(self):
      tf_proto, tf_text = jax_to_ir.jax_to_tf(
          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 tf debug txt contains a broadcast, add, and multiply.
      self.assertIn("BroadcastTo", tf_text)
      self.assertIn("AddV2", tf_text)
      self.assertIn("Mul", tf_text)
  
      # Check that we can re-import our graphdef.
      gdef = tf.compat.v1.GraphDef()
      gdef.ParseFromString(tf_proto)
      g = tf.Graph()
      with g.as_default():
          tf.import_graph_def(gdef, name="")
  
      # Check that the HLO parameters are named as we specified.
      ops = {
          o.name: o
          for o in g.get_operations()
          if o.name in ("y", "a", "x", "jax2tf_out")
      }
      self.assertLen(ops, 4)
      self.assertIdentityOp(ops["y"], [128], jnp.float32)
      self.assertIdentityOp(ops["a"], [], jnp.float32)
      self.assertIdentityOp(ops["x"], [128, 2], jnp.float32)
      self.assertIdentityOp(ops["jax2tf_out"], [128, 2], jnp.float32)
  
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 4 for calling the JaxToIRTest obj's underlying member method code with example usage: obj.test_jax_to_tf_axpy() without return types