code: def test_compile(self, golden):
      mod = golden.create_module()
      golden.create_all_variables(mod)
  
      @tf.function
      def forward():
          f = lambda: golden.forward(mod)
          out = tf.xla.experimental.compile(f)
          if len(out) == 1:
              return out[0]
          else:
              return out if out else None
  
      if self.primary_device == "TPU":
          # TODO(b/132329316) Remove when `xla.compile` allows tf.device(TPU).
          with tf.device(None):
              xla_out = forward()
          atol = golden.tpu_atol
      else:
          xla_out = forward()
          atol = 1e-3
  
      if golden.deterministic and not golden.has_side_effects:
          out = golden.forward(mod)
          tree.map_structure(
              functools.partial(self.assertAllClose, atol=atol),
              out,
              xla_out,
          )
  
file_path: sonnet/sonnet/src/conformance/xla_test.py
imports: import functools
  from absl.testing import parameterized
  from sonnet.src import test_utils
  from sonnet.src.conformance import goldens
  import tensorflow
  import tree
label_desc: Use image node_id 1 for calling the XLATest obj's underlying member method code with example usage: obj.test_compile(golden) and returns: out, unknown