| 213 | @test_util.run_gpu_only |
| 214 | @reset_eager |
| 215 | def testJit(self): |
| 216 | self.assertEqual(config.get_optimizer_jit(), False) |
| 217 | |
| 218 | # the following function should cause Op fusion to occur. However, there is |
| 219 | # unfortunately no straightforward way to ensure this. We will just have to |
| 220 | # settle for creating a test that can trigger JIT. |
| 221 | @def_function.function |
| 222 | def fun(a, b): |
| 223 | c = a * b |
| 224 | d = c + a |
| 225 | return d |
| 226 | |
| 227 | a = constant_op.constant([2., 2.]) |
| 228 | b = constant_op.constant([2., 2.]) |
| 229 | |
| 230 | self.evaluate(fun(a, b)) |
| 231 | |
| 232 | config.set_optimizer_jit(True) |
| 233 | self.assertEqual(config.get_optimizer_jit(), True) |
| 234 | self.assertEqual(config.get_optimizer_jit(), |
| 235 | context.context().optimizer_jit) |
| 236 | |
| 237 | self.evaluate(fun(a, b)) |
| 238 | |
| 239 | config.set_optimizer_jit(False) |
| 240 | self.assertEqual(config.get_optimizer_jit(), False) |
| 241 | self.assertEqual(config.get_optimizer_jit(), |
| 242 | context.context().optimizer_jit) |
| 243 | |
| 244 | self.evaluate(fun(a, b)) |
| 245 | |
| 246 | @parameterized.named_parameters( |
| 247 | ('LayoutOptimizer', 'layout_optimizer'), |