(self)
| 992 | ) |
| 993 | |
| 994 | def test_remat(self): |
| 995 | num_iters = 3 |
| 996 | |
| 997 | with self._dummy_context(): |
| 998 | ref_carry, ref_ys = self._invoke(num_iters=num_iters, xs={}) |
| 999 | |
| 1000 | with self._dummy_context(): |
| 1001 | test_carry, test_ys = self._invoke( |
| 1002 | num_iters=num_iters, |
| 1003 | xs={}, |
| 1004 | remat_kwargs=dict(policy=jax.checkpoint_policies.everything_saveable), |
| 1005 | ) |
| 1006 | self.assertNestedEqual(ref_carry, test_carry) |
| 1007 | self.assertNestedEqual(ref_ys, test_ys) |
| 1008 | |
| 1009 | # prevent_cse=True raises ValueError. |
| 1010 | with self._dummy_context(), self.assertRaises(ValueError): |
| 1011 | _ = self._invoke( |
| 1012 | num_iters=num_iters, |
| 1013 | xs={}, |
| 1014 | remat_kwargs=dict(prevent_cse=True), |
| 1015 | ) |
| 1016 | |
| 1017 | def test_merge_summaries_mixed_summary_and_tensor(self): |
| 1018 | """`merge_summaries=True` must accumulate WeightedSummary via .accumulate() |
nothing calls this directly
no test coverage detected