code: def test_tpu_rewrite_uses_xla_einsum(self):
      with ops.Graph().as_default() as g:
          tpu.rewrite(do_einsum)
          self.assertTrue(find_einsum(g) or find_xla_einsum(g))
  
imports: from tensorflow.python.eager import def_function
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import importer
  from tensorflow.python.framework import ops
  from tensorflow.python.layers import convolutional
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import control_flow_ops
  from tensorflow.python.ops import control_flow_util
  from tensorflow.python.ops import init_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import special_math_ops
  from tensorflow.python.ops import variable_scope
  from tensorflow.python.platform import test
  from tensorflow.python.tpu import tpu
  from tensorflow.python.tpu import tpu_feed
  from tensorflow.python.tpu import tpu_replication
  from tensorflow.python.tpu import training_loop
  from tensorflow.python.tpu.ops import tpu_ops
label_desc: Use image node_id 1 for calling the TPUXlaEinsumTest obj's underlying member method code with example usage: obj.test_tpu_rewrite_uses_xla_einsum() without return types