(self, *args, **kwargs)
| 1114 | """Holds a map from replica to TPU variables whose values are kept in sync.""" |
| 1115 | |
| 1116 | def _assign_func(self, *args, **kwargs): |
| 1117 | with _enter_or_assert_strategy(self._distribute_strategy): |
| 1118 | if (distribution_strategy_context.in_cross_replica_context() |
| 1119 | and (_enclosing_tpu_context() is not None)): |
| 1120 | f = kwargs.pop("f") |
| 1121 | return self._distribute_strategy.extended.update( |
| 1122 | self, f, args=args, kwargs=kwargs) |
| 1123 | else: |
| 1124 | return MirroredVariable._assign_func(self, *args, **kwargs) |
| 1125 | |
| 1126 | def assign_sub(self, *args, **kwargs): |
| 1127 | assign_sub_fn = _make_raw_assign_fn( |
no test coverage detected