def __enter__(self):
"""Runs ops in parallel, makes variables which save independent buffers."""
if (self._device_scope is not None or self._saving_scope is not None):
raise AssertionError(
"Re-entered a ParallelDevice scope without first exiting it.")
self._device_scope = ops.device(self._name)
self._saving_scope = saving.independent_buffers(self)
self._device_scope.__enter__()
# TODO(allenl): Fixing saving in Python is a bit odd. One alternative would
# be to provide a hook for the custom device to create save specs/etc., then
# call that hook from the default variable implementation if the variable is
# on a custom device. We'll likely want similar hooks for repr() and such.
self._saving_scope.__enter__()
return self
def scope(self):
"""Runs ops in parallel, makes variables which save independent buffers."""
with ops.device(self.name), saving.independent_buffers(self):
yield