Test that control_flow_util can check that we're in a TPU context.
(self)
| 42 | |
| 43 | @test_util.deprecated_graph_mode_only |
| 44 | def testIsInContext(self): |
| 45 | """Test that control_flow_util can check that we're in a TPU context.""" |
| 46 | z1 = array_ops.identity(1) |
| 47 | pivot = control_flow_ops.no_op() |
| 48 | context = tpu.TPUReplicateContext(b"context", 1, pivot=pivot) |
| 49 | context.Enter() |
| 50 | z2 = array_ops.identity(1) |
| 51 | context.Exit() |
| 52 | self.assertFalse(control_flow_util.IsInXLAContext(z1.op)) |
| 53 | self.assertTrue(control_flow_util.IsInXLAContext(z2.op)) |
| 54 | |
| 55 | |
| 56 | class TPULayerRewriteTest(test.TestCase): |