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)
