def assign(self, var, *args, **kwargs): if tpu_util.enclosing_tpu_context() is None: return super(TPUOnReadPolicy, self).assign(var, *args, **kwargs) else: return tpu_util.make_raw_assign_fn( gen_resource_variable_ops.assign_variable_op)(var, *args, **kwargs)
def assign(self, *args, **kwargs): if tpu_util.enclosing_tpu_context() is None: return values.SyncOnReadVariable.assign(self, *args, **kwargs) else: return tpu_util.make_raw_assign_fn( gen_resource_variable_ops.assign_variable_op)(self, *args, **kwargs)
def assign(self, value, use_locking=False, name=None, read_value=True): tpu_context = tpu_util.enclosing_tpu_context() if (self._is_replicated_or_sharded_to_logical_cores() and tpu_context is None): assign_fn = lambda v, *a, **ka: v.assign(*a, **ka) return self._update(update_fn=assign_fn, value=value, use_locking=use_locking, name=name, read_value=read_value) if (tpu_util.enclosing_tpu_context() and self.aggregation == variable_scope.VariableAggregation.NONE): return tpu_util.make_raw_assign_fn( gen_resource_variable_ops.assign_variable_op)( self, value=value, use_locking=use_locking, name=name, read_value=read_value) return assign(self, value, use_locking=use_locking, name=name, read_value=read_value)
def assign(var, value, use_locking=False, name=None, read_value=True): assign_fn = tpu_util.make_raw_assign_fn( gen_resource_variable_ops.assign_variable_op) return var._update( # pylint: disable=protected-access update_fn=assign_fn, value=value, use_locking=use_locking, name=name, read_value=read_value)
def assign_add(self, value, use_locking=False, name=None, read_value=True): if tpu_util.enclosing_tpu_context( ) is None or context.executing_eagerly(): assign_add_fn = lambda var, *a, **ka: var.assign_add(*a, **ka) return self._update(assign_add_fn, value=value, use_locking=use_locking, name=name, read_value=read_value) else: return tpu_util.make_raw_assign_fn( gen_resource_variable_ops.assign_add_variable_op)( self, value=value, use_locking=use_locking, name=name, read_value=read_value)
def assign(self, var, value, use_locking=False, name=None, read_value=True): if (tpu_util.enclosing_tpu_context() and var.aggregation == variable_scope.VariableAggregation.NONE): return tpu_util.make_raw_assign_fn( gen_resource_variable_ops.assign_variable_op)( var, value=value, use_locking=use_locking, name=name, read_value=read_value) return assign(var, value, use_locking=use_locking, name=name, read_value=read_value)