"""Allows attribute assignment to `Trackable` objects with no dependency.

__slots__ = ["value"]

def __init__(self, value):

"""Determine if a tuple has any trackable components."""

# pylint: disable=unidiomatic-typecheck

# Exact type checking to avoid mucking up custom logic in list/dict

# subclasses, e.g. collections.Counter.

for element in t:

if isinstance(element, NoDependency):

return True # We should remove the NoDependency object from the tuple.

if isinstance(element, base.Trackable):

return True

if type(element) == dict:

return True

if type(element) == collections.OrderedDict:

return True

if type(element) == list:

return True

if isinstance(element, tuple) and _should_wrap_tuple(element):

return True

# There are no trackable elements or data structures. Tuples are immutable, so

# mutation isn't a concern. Don't wrap.

return False

"""Wraps input value into trackable data structures.

# pylint: disable=unidiomatic-typecheck

# Exact type checking to avoid mucking up custom logic in list/dict

# subclasses, e.g. collections.Counter.

if isinstance(value, NoDependency):

return value.value

if isinstance(value, base.Trackable):

return value # Skip conversion for already trackable objects.

elif type(value) == dict:

return _DictWrapper(value)

elif type(value) == collections.OrderedDict:

return _DictWrapper(value)

elif type(value) == list:

return ListWrapper(value)

elif isinstance(value, tuple) and _should_wrap_tuple(value):

# There are trackable elements or data structures. Wrap the tuple.

return _TupleWrapper(value)

else:

return value

"""Adds dependencies, generally called from __setattr__.

if isinstance(value, NoDependency):

add_dependency = False

else:

add_dependency = True

value = wrap_or_unwrap(value)

if not add_dependency:

return value

if isinstance(value, base.Trackable):

trackable._track_trackable( # pylint: disable=protected-access

value, name=name,

# Allow the user to switch the Trackable which is tracked by this

# name, since assigning a new variable to an attribute has

# historically been fine (e.g. Adam did this).

overwrite=True)

def __init__(self, value): # pylint: disable=super-init-not-called

self._value = value

def __str__(self):

return ("Only trackable objects (such as Layers or Optimizers) may be "

f"stored in a List object. Got {self._value}, which does not "

"""Base class for data structures which contain trackable objects."""

def __init__(self):

# Attributes prefixed with "_self_" for compatibility with

# wrapt.ObjectProxy. All additional attrs MUST conform to this pattern, as

# extending `__slots__` on a subclass of ObjectProxy breaks in a variety of

# ways.

self._self_trainable = True

self._self_extra_variables = []

self._self_attribute_sentinel = layer_utils.AttributeSentinel(True)

@property

def _attribute_sentinel(self):

return self._self_attribute_sentinel

@property

def trainable(self):

return self._self_trainable

@trainable.setter

def trainable(self, value):

self._self_trainable = value

def _track_value(self, value, name):

"""Add a dependency on `value`."""

value = sticky_attribute_assignment(

trackable=self, value=value, name=name)

if isinstance(value, variables.Variable):

self._self_extra_variables.append(value)

if not isinstance(value, base.Trackable):

raise _UntrackableError(value)

if hasattr(value, "_use_resource_variables"):

# In subclassed models, legacy layers (tf.layers) must always use

# resource variables.

value._use_resource_variables = True # pylint: disable=protected-access

value_attribute_sentinel = getattr(value, "_attribute_sentinel", None)

if value_attribute_sentinel:

value_attribute_sentinel.add_parent(self._attribute_sentinel)

return value

@property

def _values(self):

"""An iterable/sequence which may contain trackable objects."""

raise NotImplementedError("Abstract method")

@property

def _layers(self):

"""All Layers and Layer containers, including empty containers."""

# Filter objects on demand so that wrapper objects use values from the thing

# they're wrapping if out of sync.

collected = []

for obj in self._values:

if (isinstance(obj, TrackableDataStructure)

or layer_utils.is_layer(obj)

or layer_utils.has_weights(obj)):

collected.append(obj)

return collected

@property

def layers(self):

return list(layer_utils.filter_empty_layer_containers(self._layers))

@property

def trainable_weights(self):

if not self._self_trainable:

return []

trainable_variables = []

for obj in self._values:

if isinstance(obj, base.Trackable) and hasattr(

obj, "trainable_variables"):

trainable_variables += obj.trainable_variables

trainable_extra_variables = [

v for v in self._self_extra_variables if v.trainable

]

return trainable_variables + trainable_extra_variables

@property

def non_trainable_weights(self):

trainable_extra_variables = [

v for v in self._self_extra_variables if v.trainable

]

non_trainable_extra_variables = [

v for v in self._self_extra_variables if not v.trainable

]

non_trainable_variables = []

for obj in self._values:

if isinstance(obj, base.Trackable) and hasattr(

obj, "non_trainable_variables"):

non_trainable_variables += obj.non_trainable_variables

if not self._self_trainable:

# Return order is all trainable vars, then all non-trainable vars.

trainable_variables = []

for obj in self._values:

if isinstance(obj, base.Trackable) and hasattr(

obj, "trainable_variables"):

trainable_variables += obj.trainable_variables

non_trainable_variables = (

trainable_variables + trainable_extra_variables +

non_trainable_variables + non_trainable_extra_variables)

else:

non_trainable_variables = (

non_trainable_variables + non_trainable_extra_variables)

return non_trainable_variables

@property

def weights(self):

return self.trainable_weights + self.non_trainable_weights

@property

def trainable_variables(self):

return self.trainable_weights

@property

def non_trainable_variables(self):

return self.non_trainable_weights

@property

def variables(self):

return self.weights

@property

def updates(self):

"""Aggregate updates from any `Layer` instances."""

# Updates and conditional losses are forwarded as-is rather than being

# filtered based on inputs, since this is just a container and won't ever

# have any inputs.

aggregated = []

for layer in self.layers:

if hasattr(layer, "updates"):

aggregated += layer.updates

return aggregated

@property

def losses(self):

"""Aggregate losses from any `Layer` instances."""

aggregated = []

for layer in self.layers:

if hasattr(layer, "losses"):

aggregated += layer.losses

return aggregated

def __hash__(self):

# Support object-identity hashing, so these structures can be used as keys

# in sets/dicts.

return id(self)

def __eq__(self, other):

# Similar to Tensors, trackable data structures use object-identity

# equality to support set/dict membership.

"""An append-only sequence type which is trackable.

def __init__(self, *args, **kwargs):

"""Construct a new sequence. Arguments are passed to `list()`."""

super().__init__()

self._storage = self._make_storage(*args, **kwargs)

for index, element in enumerate(self._storage):

self._storage[index] = self._track_value(

element, name=self._name_element(index))

def copy(self):

return type(self)(copy.copy(self._storage))

def __copy__(self):

return self.copy()

def __deepcopy__(self, memo):

return type(self)(copy.deepcopy(self._storage, memo))

def _make_storage(self, *args, **kwargs):

"""Determines the backing storage (overridden in subclasses)."""

return list(*args, **kwargs)

def _name_element(self, index):

return "%d" % (index,)

@property

def _values(self):

"""Collect values for TrackableDataStructure."""

return self

def append(self, value):

"""Add a new trackable value."""

value = self._track_value(value, self._name_element(len(self._storage)))

self._storage.append(value)

def extend(self, values):

"""Add a sequence of trackable values."""

for value in values:

self.append(value)

def __iadd__(self, values):

self.extend(values)

return self

def __add__(self, other):

return self._storage + getattr(other, "_storage", other)

def __imul__(self, y):

if y <= 0:

raise ValueError(

f"List only supports append, multiplying in place by {y} removes "

"elements.")

n = len(self._storage)

for _ in range(y - 1):

for i in range(n):

self.append(self._storage[i])

return self

def __mul__(self, n):

return self._storage * n

def __rmul__(self, n):

return self * n

def __radd__(self, other):

return other + self._storage

def __getitem__(self, key):

return self._storage[key]

def __getslice__(self, i, j):

return self._storage[slice(i, j)]

def __len__(self):

return len(self._storage)

def __repr__(self):

return "List(%s)" % (repr(self._storage),)

def __sizeof__(self):

List,

collections_abc.MutableSequence,

# Shadowed, but there for isinstance checks.

list):

"""Wraps the built-in `list` to support restore-on-create for variables.

def __init__(self, wrapped_list):

"""Construct a new list wrapper.

# Monotonic flags which indicate this object would not be restored properly,

# and therefore should throw an error on save to avoid giving the impression

# that restoring it will work.

self._non_append_mutation_value = False

self._external_modification_value = False

super().__init__(wrapped_list)

self._last_wrapped_list_snapshot = list(self._storage)

@property

def _non_append_mutation(self):

return self._non_append_mutation_value

@_non_append_mutation.setter

def _non_append_mutation(self, value):

# Trackable only cares that a mutation occurred at some point; when

# attempting to save it checks whether a mutation occurred and the object is

# in a "dirty" state but otherwise the specifics of how it got to that state

# are ignored. By contrast, the attribute cache needs to signal the mutation

# immediately since a caller could query the value of an attribute (And

# should not hit the cached value since the mutation may have affected the

# result.)

self._attribute_sentinel.invalidate_all()

self._non_append_mutation_value = value

@property

def _external_modification(self):

return self._external_modification_value

@_external_modification.setter

def _external_modification(self, value):

# Invalidate for the same reason as `_non_append_mutation`

self._attribute_sentinel.invalidate_all()

self._external_modification_value = value

# pylint: disable=protected-access

def __copy__(self):

copied = super().__copy__()

copied._non_append_mutation = self._non_append_mutation

copied._external_modification = self._external_modification

return copied

def __deepcopy__(self, memo):

copied = super().__deepcopy__(memo)

copied._non_append_mutation = self._non_append_mutation

copied._external_modification = self._external_modification

return copied

# pylint: enable=protected-access

def __reduce_ex__(self, protocol):

return (self.__class__,

(self._storage,))

def _make_storage(self, wrapped_list):

"""Use the user's original list for storage."""

return wrapped_list

def _check_external_modification(self):

"""Checks for any changes to the wrapped list not through the wrapper."""

if self._external_modification or self._non_append_mutation:

return

if self._storage != self._last_wrapped_list_snapshot:

self._external_modification = True

self._last_wrapped_list_snapshot = None

def _update_snapshot(self):

"""Acknowledges tracked changes to the wrapped list."""

# Mutation tracking for attributes reuses the same infrastructure as

# Trackable mutation tracking.

self._attribute_sentinel.invalidate_all()

if self._external_modification or self._non_append_mutation:

return

self._last_wrapped_list_snapshot = list(self._storage)

def _trackable_children(self, save_type=base.SaveType.CHECKPOINT, **kwargs):

self._check_external_modification()

if self._non_append_mutation:

raise ValueError(

f"Unable to save the object {self} (a list wrapper constructed to "

"track trackable TensorFlow objects). A list element was replaced "

"(__setitem__, __setslice__), deleted (__delitem__, __delslice__), "

"or moved (sort). In order to support restoration on object "

"creation, tracking is exclusively for append-only data structures."

"\n\nIf you don't need this list checkpointed, wrap it in a "

"non-trackable object; it will be subsequently ignored.")

if self._external_modification:

raise ValueError(

f"Unable to save the object {self} (a list wrapper constructed to "

"track trackable TensorFlow objects). The wrapped list was modified "

f"outside the wrapper (its final value was {self._storage}, its value"

" when a checkpoint dependency was added was "

f"{self._last_wrapped_list_snapshot}), which breaks "

"restoration on object creation.\n\nIf you don't need this list "

"checkpointed, wrap it in a NoDependency object; it will be "

"subsequently ignored.")

children = super()._trackable_children(save_type, **kwargs)

if save_type == base.SaveType.SAVEDMODEL:

# Add functions to be serialized.

children.update({

str(key): value

for key, value in enumerate(self)

if _is_function(value)

})

return children

def _has_mutation_or_trackable(self):

"""Short-circuits a check for trackables if there's already a mutation."""

if self._non_append_mutation:

return True

return any(isinstance(element, base.Trackable) for element in self._storage)

def __delitem__(self, key):

self._check_external_modification()

if self._has_mutation_or_trackable():

self._non_append_mutation = True

del self._storage[key]

self._update_snapshot()

def __setitem__(self, key, value):

self._check_external_modification()

if isinstance(key, slice):

# Note: this is quite inefficient, but the list API supports a broad range

# of slice setters (e.g. truncate, extend, replace) and imitating this

# for a range of Python versions is non-trivial.

storage_copy = list(self._storage)

self._storage[key] = value

len_before = len(storage_copy)

len_now = len(self._storage)

for i in range(max(len_before, len_now)):

value_now = self._storage[i] if i < len_now else None

value_before = storage_copy[i] if i < len_before else None

if isinstance(value_before, base.Trackable):

self._non_append_mutation = True

if value_now is not None and value_now != value_before:

self._storage[i] = self._track_value(self._storage[i],

self._name_element(i))

else:

if isinstance(self._storage[key], base.Trackable):

self._non_append_mutation = True

self._storage[key] = self._track_value(value, self._name_element(key))

self._update_snapshot()

def append(self, value):

"""Add a new trackable value."""

self._check_external_modification()

super().append(value)

self._update_snapshot()

def extend(self, values):

"""Add a sequence of trackable values."""

self._check_external_modification()

super().extend(values)

self._update_snapshot()

def __imul__(self, y):

if y <= 0:

self._check_external_modification()

if self._has_mutation_or_trackable():

self._non_append_mutation = True

self._storage *= y

self._update_snapshot()

return self

# Relies on super() calling append, which updates the snapshot.

return super().__imul__(y)

def __eq__(self, other):

return self._storage == getattr(other, "_storage", other)

def __ne__(self, other):

return self._storage != getattr(other, "_storage", other)

def __lt__(self, other):

return self._storage < getattr(other, "_storage", other)

def __le__(self, other):

return self._storage <= getattr(other, "_storage", other)

def __gt__(self, other):

return self._storage > getattr(other, "_storage", other)

def __ge__(self, other):

return self._storage >= getattr(other, "_storage", other)

def __hash__(self):

# List wrappers need to compare like regular lists, and so like regular

# lists they don't belong in hash tables.

raise TypeError("unhashable type: 'ListWrapper'")

def insert(self, index, obj):

self._check_external_modification()

if (self._has_mutation_or_trackable() or isinstance(obj, base.Trackable)):

self._non_append_mutation = True

self._storage.insert(index, obj)

self._update_snapshot()

def sort(self):

self._check_external_modification()

if self._has_mutation_or_trackable():

self._non_append_mutation = True

self._storage.sort()

self._update_snapshot()

def __setslice__(self, i, j, y):

self.__setitem__(slice(i, j), y)

def __delslice__(self, i, j):

self._check_external_modification()

if self._has_mutation_or_trackable():

self._non_append_mutation = True

del self._storage[slice(i, j)]

self._update_snapshot()

def _track_value(self, value, name):

"""Allows storage of non-trackable objects."""

try:

value = super()._track_value(value=value, name=name)

except ValueError:

# Even if this value isn't trackable, we need to make sure

# NoDependency objects get unwrapped.

value = sticky_attribute_assignment(

trackable=self, value=value, name=name)

return value

def __repr__(self):

"""An append-only trackable mapping data structure with string keys.

def __init__(self, *args, **kwargs):

"""Construct a new sequence. Arguments are passed to `dict()`."""

super().__init__()

self._storage = self._make_storage(*args, **kwargs)

self._storage.update(

{key: self._track_value(

value, name=self._name_element(key))

for key, value in self._storage.items()})

def __copy__(self):

return type(self)(copy.copy(self._storage))

def __deepcopy__(self, memo):

return type(self)(copy.deepcopy(self._storage, memo))

def _make_storage(self, *args, **kwargs):

return dict(*args, **kwargs)

@property

def _values(self):

"""Collect values for TrackableDataStructure."""

# Sort items deterministically by key

ordered = list(zip(*sorted(self.items(), key=lambda it: it[0])))

if ordered:

return ordered[1]

return []

def _name_element(self, key):

if not isinstance(key, str):

raise TypeError(

f"Mapping accepts only string keys, but got a key {repr(key)}.")

return str(key)

def __setitem__(self, key, value):

name = self._name_element(key)

value = self._track_value(value, name=name)

current_value = self._storage.setdefault(key, value)

if current_value is not value:

raise ValueError(

"Mappings are an append-only data structure. Tried to overwrite the "

f"key '{key}' with value {value}, but it already contains "

f"{current_value}")

def update(self, *args, **kwargs):

for key, value in dict(*args, **kwargs).items():

self[key] = value

def __getitem__(self, key):

return self._storage[key]

def __len__(self):

return len(self._storage)

def __repr__(self):

return "Mapping(%s)" % (repr(self._storage),)

def __iter__(self):

"""Wraps built-in dicts to support restore-on-create for variables.

def __init__(self, wrapped_dict=None):

if wrapped_dict is None:

# Allow zero-argument construction, e.g. from session.run's re-wrapping.

wrapped_dict = {}

if not isinstance(wrapped_dict, collections_abc.Mapping):

# Allow construction from a sequence, e.g. from nest.pack_sequence_as.

wrapped_dict = dict(wrapped_dict)

wrapt.ObjectProxy.__init__(self, wrapped_dict)

TrackableDataStructure.__init__(self)

self._self_non_string_key = False

self._self_external_modification = False

self.__wrapped__.update(

{key: self._track_value(

value, name=self._name_element(key))

for key, value in self.__wrapped__.items()})

self._update_snapshot()

def __reduce_ex__(self, protocol):

return (self.__class__,

(self.__wrapped__,))

def __getattribute__(self, name):

if (hasattr(type(self), name)

and isinstance(getattr(type(self), name), property)):

# Bypass ObjectProxy for properties. Whether this workaround is necessary

# appears to depend on the Python version but not the wrapt version: 3.4

# in particular seems to look up properties on the wrapped object instead

# of the wrapper without this logic.

return object.__getattribute__(self, name)

else:

return super().__getattribute__(name)

def copy(self):

return copy.copy(self)

# pylint: disable=protected-access

def __copy__(self):

copied = _DictWrapper(copy.copy(self.__wrapped__))

copied._self_external_modification = self._self_external_modification

copied._self_non_string_key = self._self_non_string_key

return copied

def __deepcopy__(self, memo):

copied = _DictWrapper(copy.deepcopy(self.__wrapped__, memo))

copied._self_external_modification = self._self_external_modification

copied._self_non_string_key = self._self_non_string_key

return copied

# pylint: enable=protected-access

@property

def _values(self):

"""Collect values for TrackableDataStructure."""

# Sort items deterministically by key

ordered = list(zip(*sorted(self.items(), key=lambda it: it[0])))

if ordered:

return ordered[1]

return []

def _trackable_children(self, save_type=base.SaveType.CHECKPOINT, **kwargs):

"""Check that the object is saveable before listing its dependencies."""

self._check_self_external_modification()

if self._self_non_string_key:

raise ValueError(

f"Unable to save the object {self} (a dictionary wrapper constructed "

"automatically on attribute assignment). The wrapped dictionary "

"contains a non-string key which maps to a trackable object or "

"mutable data structure.\n\nIf you don't need this dictionary "

"checkpointed, wrap it in a non-trackable "

"object; it will be subsequently ignored.")

if self._self_external_modification:

raise ValueError(

f"Unable to save the object {self} (a dictionary wrapper constructed "

"automatically on attribute assignment). The wrapped dictionary was "

f"modified outside the wrapper (its final value was {self}, its value"

" when a checkpoint dependency was added was "

f"{self._self_last_wrapped_dict_snapshot}), which breaks "

"restoration on object creation.\n\nIf you don't need this "

"dictionary checkpointed, wrap it in a "

"non-trackable object; it will be subsequently ignored.")

assert not self._dirty # Any reason for dirtiness should have an exception.

children = super()._trackable_children(save_type, **kwargs)

if save_type == base.SaveType.SAVEDMODEL:

# Add functions to be serialized.

children.update(

{key: value for key, value in self.items() if _is_function(value)})

return children

@property

def _dirty(self):

"""Check if there has already been a mutation which prevents saving."""

return (self._self_external_modification

or self._self_non_string_key)

def _check_self_external_modification(self):

"""Checks for any changes to the wrapped dict not through the wrapper."""

if self._dirty:

return

if self != self._self_last_wrapped_dict_snapshot:

self._self_external_modification = True

self._self_last_wrapped_dict_snapshot = None

def _update_snapshot(self):

"""Acknowledges tracked changes to the wrapped dict."""

self._attribute_sentinel.invalidate_all()

if self._dirty:

return

self._self_last_wrapped_dict_snapshot = dict(self)

def _track_value(self, value, name):

"""Allows storage of non-trackable objects."""

if isinstance(name, str):

string_key = True

else:

name = "-non_string_key"

string_key = False

try:

no_dependency = isinstance(value, NoDependency)

value = super()._track_value(value=value, name=name)

if not (string_key or no_dependency):

# A non-string key maps to a trackable value. This data structure

# is not saveable.

self._self_non_string_key = True

return value

except ValueError:

# Even if this value isn't trackable, we need to make sure

# NoDependency objects get unwrapped.

return sticky_attribute_assignment(

trackable=self, value=value, name=name)

def _name_element(self, key):

"""Tells TrackableDataStructure to use keys as names as-is."""

return key

def __setitem__(self, key, value):

"""Allow any modifications, but possibly mark the wrapper as unsaveable."""

self._check_self_external_modification()

self._maybe_initialize_trackable()

no_dep = isinstance(value, NoDependency)

if isinstance(key, str):

value = self._track_value(value, name=key)

else:

value = wrap_or_unwrap(value)

if not no_dep and isinstance(value, base.Trackable):

# Non-string keys are OK as long as we have no reason to add a

# dependency on the value (either because the value is not

# trackable, or because it was wrapped in a NoDependency object).

self._self_non_string_key = True

self.__wrapped__[key] = value

self._update_snapshot()

def __delitem__(self, key):

self._check_self_external_modification()

del self.__wrapped__[key]

self._update_snapshot()

def __repr__(self):

return "DictWrapper(%s)" % (repr(self.__wrapped__),)

def __hash__(self):

raise TypeError("unhashable type: 'DictWrapper'")

def __eq__(self, other):

# Override the TrackableDataStructure "== -> is" forwarding and go back to

# the wrapt implementation.

return self.__wrapped__ == other

def update(self, *args, **kwargs):

for key, value in dict(*args, **kwargs).items():

"""Trackable wrapper for tuples and namedtuples."""

def __init__(self, original_wrapped_tuple=()):

add_dependency = []

substituted_wrapped_tuple = []

for element in original_wrapped_tuple:

if isinstance(element, NoDependency):

add_dependency.append(False)

else:

add_dependency.append(True)

substituted_wrapped_tuple.append(wrap_or_unwrap(element))

try:

fields = original_wrapped_tuple._fields

except AttributeError:

# Not a namedtuple

is_namedtuple = False

else:

is_namedtuple = True

original_type = type(original_wrapped_tuple)

# Flag to poison saving if we can't re-construct a namedtupled because its

# __new__ takes different keyword arguments than its _fields.

self._self_tuple_is_constructable = True

if is_namedtuple:

try:

# NamedTuples take N arguments, unlike tuple which takes a sequence.

substituted_wrapped_tuple = original_type(

**dict(zip(fields, substituted_wrapped_tuple)))

except TypeError:

wrapt.ObjectProxy.__init__(self, original_wrapped_tuple)

TrackableDataStructure.__init__(self)

self._self_tuple_is_constructable = False

return

else:

substituted_wrapped_tuple = original_type(substituted_wrapped_tuple)

wrapt.ObjectProxy.__init__(self, substituted_wrapped_tuple)