Calls model_fn and verifies it is compiled.
(self, model_fn, params)
| 79 | }), |
| 80 | ) |
| 81 | def test_compile(self, model_fn, params): |
| 82 | """Calls model_fn and verifies it is compiled.""" |
| 83 | with test.mock.patch.object(xla, 'compile') as mock_xla_compile: |
| 84 | loss = constant_op.constant(_EXPECTED_LOSS) |
| 85 | mock_xla_compile.return_value = [loss] |
| 86 | |
| 87 | features, labels = make_dummy_features_labels() |
| 88 | estimator_spec = model_fn( |
| 89 | features=features, labels=labels, mode=_TRAIN, params=params or {}) |
| 90 | |
| 91 | self.assertEqual(mock_xla_compile.call_count, 1) |
| 92 | self.assertEqual(estimator_spec.mode, _TRAIN) |
| 93 | |
| 94 | with self.test_session() as sess: |
| 95 | self.assertEqual(sess.run(estimator_spec.loss), sess.run(loss)) |
| 96 | self.assertEqual(sess.run(estimator_spec.train_op), sess.run(loss)) |
| 97 | |
| 98 | @parameterized.named_parameters( |
| 99 | ('test_use_tpu_true_hparams', decorated_model_fn, |
nothing calls this directly
no test coverage detected