MCPcopy Create free account
hub / github.com/DeepRec-AI/DeepRec / _assign_func

Method _assign_func

tensorflow/python/distribute/values.py:1492–1521  ·  view source on GitHub ↗
(self, *args, **kwargs)

Source from the content-addressed store, hash-verified

1490 return getattr(self._v, name)
1491
1492 def _assign_func(self, *args, **kwargs):
1493 with _enter_or_assert_strategy(self._distribute_strategy):
1494 f = kwargs.pop("f")
1495 if distribution_strategy_context.in_cross_replica_context():
1496 update_device = distribute_lib.get_update_device()
1497 if update_device is not None:
1498 # We are calling an assign function in an update context.
1499 return f(self._v, *args, **kwargs)
1500
1501 # We are calling an assign function in cross replica context, wrap it in
1502 # an update call.
1503 return self._distribute_strategy.extended.update(
1504 self, f, args=args, kwargs=kwargs)
1505 else:
1506 replica_context = distribution_strategy_context.get_replica_context()
1507 assert replica_context
1508 # We are calling an assign function in replica context.
1509 # We reduce the value we want to assign/add/sub. More details about how
1510 # we handle the different use cases can be found in the _reduce method.
1511 # We call the function with the reduced value.
1512 if self._aggregation == vs.VariableAggregation.NONE:
1513 raise ValueError(_aggregation_error_msg.format(
1514 variable_type="AggregatingVariable"))
1515
1516 def merge_fn(strategy, value, *other_args, **other_kwargs):
1517 v = _apply_aggregation(strategy, value, self._aggregation, self)
1518 return strategy.extended.update(
1519 self, f, args=(v,) + other_args, kwargs=other_kwargs)
1520
1521 return replica_context.merge_call(merge_fn, args=args, kwargs=kwargs)
1522
1523 def assign_sub(self, *args, **kwargs):
1524 assign_sub_fn = lambda var, *a, **kw: var.assign_sub(*a, **kw)

Callers 3

assign_subMethod · 0.95
assign_addMethod · 0.95
assignMethod · 0.95

Calls 6

fFunction · 0.50
popMethod · 0.45
updateMethod · 0.45
formatMethod · 0.45
merge_callMethod · 0.45

Tested by

no test coverage detected