def build(self, _): for column in self._feature_columns: with variable_scope._pure_variable_scope( # pylint: disable=protected-access self.name, partitioner=self._partitioner): with variable_scope._pure_variable_scope( # pylint: disable=protected-access feature_column_v2. _sanitize_column_name_for_variable_scope( # pylint: disable=protected-access column.name)): column.create_state(self._state_manager) super(_BaseFeaturesLayer, self).build(None)
def wrapper(obj, *args, **kwargs): def default_context_manager(reuse=None): variable_scope = obj.variable_scope return tf.variable_scope(variable_scope, reuse=reuse) variable_scope_context_manager = getattr(obj, "_enter_variable_scope", default_context_manager) graph = tf.get_default_graph() # Temporarily enter the variable scope to capture it with variable_scope_context_manager() as tmp_variable_scope: variable_scope = tmp_variable_scope with variable_scope_ops._pure_variable_scope( variable_scope, reuse=tf.AUTO_REUSE) as pure_variable_scope: name_scope = variable_scope.original_name_scope if name_scope[-1] != "/": name_scope += "/" with tf.name_scope(name_scope): sub_scope = snt_util.to_snake_case(method.__name__) with tf.name_scope(sub_scope) as scope: out_ops = method(obj, *args, **kwargs) return out_ops
def wrapper(obj, *args, **kwargs): def default_context_manager(reuse=None): variable_scope = obj.variable_scope return tf.variable_scope(variable_scope, reuse=reuse) variable_scope_context_manager = getattr(obj, "_enter_variable_scope", default_context_manager) graph = tf.get_default_graph() # Temporarily enter the variable scope to capture it with variable_scope_context_manager() as tmp_variable_scope: variable_scope = tmp_variable_scope with variable_scope_ops._pure_variable_scope( variable_scope, reuse=tf.AUTO_REUSE) as pure_variable_scope: name_scope = variable_scope.original_name_scope if name_scope[-1] != "/": name_scope += "/" with tf.name_scope(name_scope): sub_scope = snt_util.to_snake_case(method.__name__) with tf.name_scope(sub_scope) as scope: out_ops = method(obj, *args, **kwargs) return out_ops
def __init__(self, name, func, create_scope_now=False, unique_name=None, custom_getter=None, create_graph_function=False): """Creates a template for the given function. Args: name: A name for the scope created by this template. The name will be made unique by appending `_N` to the it (see how `tf.variable_scope` treats the `default_name` for details). func: The function to apply each time. create_scope_now: Whether to create the scope at Template construction time, rather than first call. Defaults to false. Creating the scope at construction time may be more convenient if the template is to passed through much lower level code, and you want to be sure of the scope name without knowing exactly where it will be first called. If set to True, the scope will be created in the constructor, and all subsequent times in `__call__`, leading to a trailing numeral being added to the names of all created Tensors. If set to False, the scope will be created at the first call location. unique_name: When used, it overrides `name` and is not made unique. If a template of the same scope/unique_name already exists and reuse is false, an error is raised. Defaults to None. custom_getter: optional custom getter to pass to `variable_scope()` create_graph_function: When True, the first invocation of the template will execute `func` as is, to allow for variable creation; however, the second invocation and every invocation thereafter will execute `func` as a graph function. Enabling this flag gives the caller access to graph-function semantics, i.e., accesses to variables are totally ordered and side-effecting ops are not pruned. Raises: ValueError: if the name is None. """ self._func = func self._stacktrace = traceback.format_stack()[:-2] self._name = name self._unique_name = unique_name self._custom_getter = custom_getter if name is None: raise ValueError("name cannot be None.") if create_scope_now: with variable_scope._pure_variable_scope( # pylint:disable=protected-access (self._unique_name or variable_scope._get_unique_variable_scope(self._name)), # pylint:disable=protected-access custom_getter=self._custom_getter) as vs: self._variable_scope = vs else: self._variable_scope = None # This variable keeps track of whether the template has been called yet, # which is not the same as whether the scope has been created. self._variables_created = False self._create_graph_function = create_graph_function
def __init__(self, name, func, create_scope_now=False, unique_name=None, custom_getter=None, create_graph_function=False): """Creates a template for the given function. Args: name: A name for the scope created by this template. The name will be made unique by appending `_N` to the it (see how `tf.variable_scope` treats the `default_name` for details). func: The function to apply each time. create_scope_now: Whether to create the scope at Template construction time, rather than first call. Defaults to false. Creating the scope at construction time may be more convenient if the template is to passed through much lower level code, and you want to be sure of the scope name without knowing exactly where it will be first called. If set to True, the scope will be created in the constructor, and all subsequent times in `__call__`, leading to a trailing numeral being added to the names of all created Tensors. If set to False, the scope will be created at the first call location. unique_name: When used, it overrides `name` and is not made unique. If a template of the same scope/unique_name already exists and reuse is false, an error is raised. Defaults to None. custom_getter: optional custom getter to pass to `variable_scope()` create_graph_function: When True, `func` will be executed as a graph function. Enabling this flag gives the caller access to graph-function semantics, i.e., accesses to variables are totally ordered and side-effecting ops are not pruned. Raises: ValueError: if `name` is None. """ if create_graph_function: self._func = function.defun(func) else: self._func = func self._stacktrace = traceback.format_stack()[:-2] self._name = name self._unique_name = unique_name self._custom_getter = custom_getter if name is None: raise ValueError("name cannot be None.") if create_scope_now: with variable_scope._pure_variable_scope( # pylint:disable=protected-access (self._unique_name or variable_scope._get_unique_variable_scope(self._name)), # pylint:disable=protected-access custom_getter=self._custom_getter) as vs: self._variable_scope = vs else: self._variable_scope = None # This variable keeps track of whether the template has been called to # completion, which is not the same as whether the scope has been created. self._variables_created = False # `MirroredStrategy` builds the graph with multiple threads. If a # `merge_call` happens within a template, multiple calls may be in progress # simultaneously. This variable keeps track of whether any call of the # template has started. self._first_call = True
def reopen_variable_scope(var_scope, reuse=None, initializer=None, regularizer=None, caching_device=None, partitioner=None, custom_getter=None, dtype=tf.float32): """Reopen the specified `var_scope` and its `original_name_scope`. `tf.variable_scope` will not open the original name scope, even if a stored `tf.VariableScope` instance is specified. This method thus allows to open exactly the same name scope as the original one. Parameters ---------- var_scope : tf.VariableScope The variable scope instance. reuse : None | bool Whether or not to reuse the variables in opened scope? initializer, regularizer, caching_device, partitioner, custom_getter, dtype Other parameters for opening the variable scope. """ if not isinstance(var_scope, tf.VariableScope): raise TypeError('`var_scope` is expected to be an instance of ' '`tf.VariableScope`.') old_name_scope = var_scope.original_name_scope with variable_scope_ops._pure_variable_scope(var_scope, reuse=reuse, initializer=initializer, regularizer=regularizer, caching_device=caching_device, partitioner=partitioner, custom_getter=custom_getter, old_name_scope=old_name_scope, dtype=dtype) as vs: name_scope = old_name_scope if name_scope and not name_scope.endswith('/'): name_scope += '/' with tf.name_scope(name_scope): yield vs
def __init__(self, name, func, create_scope_now=False, unique_name=None, custom_getter=None): """Creates a template for the given function. Args: name: A name for the scope created by this template. The name will be made unique by appending `_N` to the it (see how `tf.variable_scope` treats the `default_name` for details). func: The function to apply each time. create_scope_now: Whether to create the scope at Template construction time, rather than first call. Defaults to false. Creating the scope at construction time may be more convenient if the template is to passed through much lower level code, and you want to be sure of the scope name without knowing exactly where it will be first called. If set to True, the scope will be created in the constructor, and all subsequent times in __call__, leading to a trailing numeral being added to the names of all created Tensors. If set to False, the scope will be created at the first call location. unique_name: When used, it overrides name_ and is not made unique. If a template of the same scope/unique_name already exists and reuse is false, an error is raised. Defaults to None. custom_getter: optional custom getter to pass to variable_scope() Raises: ValueError: if the name is None. """ self._func = func self._stacktrace = traceback.format_stack()[:-2] self._name = name self._unique_name = unique_name self._custom_getter = custom_getter if name is None: raise ValueError("name cannot be None.") if create_scope_now: with variable_scope._pure_variable_scope( # pylint:disable=protected-access (self._unique_name or variable_scope._get_unique_variable_scope(self._name)), # pylint:disable=protected-access custom_getter=self._custom_getter) as vs: self._variable_scope = vs else: self._variable_scope = None # This variable keeps track of whether the template has been called yet, # which is not the same as whether the scope has been created. self._variables_created = False
def root_variable_scope(**kwargs): """ Open the root variable scope and its name scope. Args: **kwargs: Named arguments for opening the root variable scope. """ # `tf.variable_scope` does not support opening the root variable scope # from empty name. It always prepend the name of current variable scope # to the front of opened variable scope. So we get the current scope, # and pretend it to be the root scope. scope = tf.get_variable_scope() old_name = scope.name try: scope._name = '' with variable_scope_ops._pure_variable_scope('', **kwargs) as vs: scope._name = old_name with tf.name_scope(None): yield vs finally: scope._name = old_name
def root_variable_scope(reuse=None, initializer=None, regularizer=None, caching_device=None, partitioner=None, custom_getter=None, dtype=tf.float32): """Open the root name and variable scope. Parameters ---------- reuse : None | bool Whether or not to reuse the variables in opened scope? initializer, regularizer, caching_device, partitioner, custom_getter, dtype Other parameters for opening the variable scope. """ # `tf.variable_scope` does not support opening the root variable scope # from empty name. It always prepend the name of current variable scope # to the front of opened variable scope. So we get the current scope, # and pretend it to be the root scope. scope = tf.get_variable_scope() old_name = scope.name try: scope._name = '' with variable_scope_ops._pure_variable_scope( '', reuse=reuse, initializer=initializer, regularizer=regularizer, caching_device=caching_device, partitioner=partitioner, custom_getter=custom_getter, old_name_scope='', dtype=dtype) as vs: scope._name = old_name with tf.name_scope(None): yield vs finally: scope._name = old_name
def reopen_variable_scope(var_scope, **kwargs): """ Reopen the specified `var_scope` and its original name scope. Unlike :func:`tf.variable_scope`, which does not open the original name scope even if a stored :class:`tf.VariableScope` instance is specified, this method opens exactly the same name scope as the original one. Args: var_scope (tf.VariableScope): The variable scope instance. **kwargs: Named arguments for opening the variable scope. """ if not isinstance(var_scope, tf.VariableScope): raise TypeError( '`var_scope` must be an instance of `tf.VariableScope`') old_name_scope = var_scope.original_name_scope with variable_scope_ops._pure_variable_scope(var_scope, **kwargs) as vs: name_scope = old_name_scope if name_scope and not name_scope.endswith('/'): name_scope += '/' # pragma: no cover with tf.name_scope(name_scope): yield vs
def call_method(method, obj, args, kwargs): """Calls `method` with a variable scope whose reuse flag is set correctly. The first time the wrapper is called it creates a `(tf.Graph, tf.VariableScope)` key and checks it for membership in `initialized_variable_scopes`. The check is `False` if and only if this is the first time the wrapper has been called with the key, otherwise the check is `True`. The result of this check is used as the `reuse` flag for entering the provided variable scope before calling `method`. Here are two examples of how to use the reuse_variables decorator. 1. Decorate an arbitrary instance method with a `variable_scope` attribute: ```python class Reusable(object): def __init__(self, name): with tf.variable_scope(None, default_name=name) as vs: self.variable_scope = vs @snt.reuse_variables def add_a(self, input_tensor): a = tf.get_variable("a", shape=input_tensor.get_shape()) return a + input_tensor obj = Reusable("reusable") x = tf.constant(5.0) out1 = obj.add_a(x) out2 = obj.add_a(x) # out1 == out2 ``` 2. Decorating a snt.AbstractModule instance method: ```python class ReusableModule(snt.AbstractModule): @snt.reuse_variables def add_a(self, input_tensor): a = tf.get_variable("a", shape=input_tensor.get_shape()) return a + input_tensor # We don't need @snt.reuse_variables here because build is wrapped by # `tf.make_template` inside `snt.AbstractModule`. def _build(self, input_tensor): b = tf.get_variable("b", shape=input_tensor.get_shape()) return b + self.add_a(input_tensor) obj = Reusable("reusable") x = tf.constant(5.0) out1 = obj(x) out2 = obj(x) # out1 == out2 ``` Args: method: The method to wrap. obj: The object instance passed to the wrapped method. args: The positional arguments (Tensors) passed to the wrapped method. kwargs: The keyword arguments passed to the wrapped method. Returns: Output of the wrapped method. Raises: ValueError: If no variable scope is provided or if `method` is a method and a variable_scope keyword argument is also provided. """ # If @reuse_variables is combined with @property, obj is passed in args # and method is still unbound at this stage. if obj is None: obj = args[0] def default_context_manager(reuse=None): variable_scope = obj.variable_scope return tf.variable_scope(variable_scope, reuse=reuse) variable_scope_context_manager = getattr(obj, "_enter_variable_scope", default_context_manager) with ops.init_scope(): # We need `init_scope` incase we're running inside a defun. In that case # what we want is information about where the function will be called not # where the function is being built. graph = tf.get_default_graph() will_call_in_eager_context = tf.executing_eagerly() if will_call_in_eager_context: initialized_variable_scopes = initialized_variable_scopes_eager else: if graph not in initialized_variable_scopes_graph: initialized_variable_scopes_graph[graph] = set() initialized_variable_scopes = initialized_variable_scopes_graph[ graph] # Temporarily enter the variable scope to capture it with variable_scope_context_manager() as tmp_variable_scope: variable_scope = tmp_variable_scope reuse = variable_scope.name in initialized_variable_scopes # Enter the pure variable scope with reuse correctly set with variable_scope_ops._pure_variable_scope( # pylint:disable=protected-access variable_scope, reuse=reuse) as pure_variable_scope: current_name_scope = tf.get_default_graph().get_name_scope() # Force tf.name_scope to treat current_name_scope as an "absolute" scope # so we can re-enter it. if current_name_scope and current_name_scope[-1] != "/": current_name_scope += "/" with tf.name_scope(current_name_scope): module_name = pure_variable_scope.name method_name = to_snake_case(method.__name__) method_name_scope = "{}/{}".format(module_name, method_name) with tf.name_scope(method_name_scope) as scope: if hasattr(obj, "_capture_variables"): with obj._capture_variables(): # pylint: disable=protected-access out_ops = method(*args, **kwargs) else: out_ops = method(*args, **kwargs) initialized_variable_scopes.add(pure_variable_scope.name) try: # If `obj` is a Sonnet module, let it know it's been connected # to the TF graph. obj._is_connected = True # pylint: disable=protected-access if not tf.executing_eagerly(): obj._add_connected_subgraph( # pylint: disable=protected-access method, out_ops, scope, *args, **kwargs) except AttributeError: pass return out_ops
def wrapper(obj, *args, **kwargs): """Calls `method` with a variable scope whose reuse flag is set correctly. The first time the wrapper is called it creates a `(tf.Graph, tf.VariableScope)` key and checks it for membership in `initialized_variable_scopes`. The check is `False` if and only if this is the first time the wrapper has been called with the key, otherwise the check is `True`. The result of this check is used as the `reuse` flag for entering the provided variable scope before calling `method`. Here are two examples of how to use the reuse_variables decorator. 1. Decorate an arbitrary instance method with a `variable_scope` attribute: ```python class Reusable(object): def __init__(self, name): with tf.variable_scope(name) as vs: self.variable_scope = vs @snt.reuse_variables def add_a(self, input_tensor): a = tf.get_variable("a", shape=input_tensor.get_shape()) return a + input_tensor obj = Reusable("reusable") x = tf.constant(5.0) out1 = obj.add_a(x) out2 = obj.add_a(x) # out1 == out2 ``` 2. Decorating a snt.AbstractModule instance method: ```python class ReusableModule(snt.AbstractModule): @snt.reuse_variables def add_a(self, input_tensor): a = tf.get_variable("a", shape=input_tensor.get_shape()) return a + input_tensor # We don't need @snt.reuse_variables here because build is wrapped by # `tf.make_template` inside `snt.AbstractModule`. def _build(self, input_tensor): b = tf.get_variable("b", shape=input_tensor.get_shape()) return b + self.add_a(input_tensor) obj = Reusable("reusable") x = tf.constant(5.0) out1 = obj(x) out2 = obj(x) # out1 == out2 ``` Args: obj: The object instance passed to the wrapped method. *args: The positional arguments (Tensors) passed to the wrapped method. **kwargs: The keyword arguments passed to the wrapped method. Returns: Output of the wrapped method. Raises: ValueError: If no variable scope is provided or if `method` is a method and a variable_scope keyword argument is also provided. """ def default_context_manager(reuse=None): variable_scope = obj.variable_scope return tf.variable_scope(variable_scope, reuse=reuse) variable_scope_context_manager = getattr(obj, "_enter_variable_scope", default_context_manager) graph = tf.get_default_graph() if graph not in initialized_variable_scopes: initialized_variable_scopes[graph] = set() initialized_variable_scopes_for_graph = initialized_variable_scopes[ graph] # Temporarily enter the variable scope to capture it with variable_scope_context_manager() as tmp_variable_scope: variable_scope = tmp_variable_scope reuse = variable_scope.name in initialized_variable_scopes_for_graph # Enter the pure variable scope with reuse correctly set with variable_scope_ops._pure_variable_scope( # pylint:disable=protected-access variable_scope, reuse=reuse) as pure_variable_scope: # Force tf.name_scope to treat variable_scope.original_name_scope as # an "absolute" scope name so we can re-enter it. name_scope = variable_scope.original_name_scope if name_scope[-1] != "/": name_scope += "/" with tf.name_scope(name_scope): sub_scope = to_snake_case(method.__name__) with tf.name_scope(sub_scope): out_ops = method(obj, *args, **kwargs) initialized_variable_scopes_for_graph.add( pure_variable_scope.name) return out_ops
def wrapper(obj, *args, **kwargs): """Calls `method` with a variable scope whose reuse flag is set correctly. The first time the wrapper is called it creates a `(tf.Graph, tf.VariableScope)` key and checks it for membership in `initialized_variable_scopes`. The check is `False` if and only if this is the first time the wrapper has been called with the key, otherwise the check is `True`. The result of this check is used as the `reuse` flag for entering the provided variable scope before calling `method`. Here are two examples of how to use the reuse_variables decorator. 1. Decorate an arbitrary instance method with a `variable_scope` attribute: ```python class Reusable(object): def __init__(self, name): with tf.variable_scope(None, default_name=name) as vs: self.variable_scope = vs @snt.reuse_variables def add_a(self, input_tensor): a = tf.get_variable("a", shape=input_tensor.get_shape()) return a + input_tensor obj = Reusable("reusable") x = tf.constant(5.0) out1 = obj.add_a(x) out2 = obj.add_a(x) # out1 == out2 ``` 2. Decorating a snt.AbstractModule instance method: ```python class ReusableModule(snt.AbstractModule): @snt.reuse_variables def add_a(self, input_tensor): a = tf.get_variable("a", shape=input_tensor.get_shape()) return a + input_tensor # We don't need @snt.reuse_variables here because build is wrapped by # `tf.make_template` inside `snt.AbstractModule`. def _build(self, input_tensor): b = tf.get_variable("b", shape=input_tensor.get_shape()) return b + self.add_a(input_tensor) obj = Reusable("reusable") x = tf.constant(5.0) out1 = obj(x) out2 = obj(x) # out1 == out2 ``` Args: obj: The object instance passed to the wrapped method. *args: The positional arguments (Tensors) passed to the wrapped method. **kwargs: The keyword arguments passed to the wrapped method. Returns: Output of the wrapped method. Raises: ValueError: If no variable scope is provided or if `method` is a method and a variable_scope keyword argument is also provided. """ def default_context_manager(reuse=None): variable_scope = obj.variable_scope return tf.variable_scope(variable_scope, reuse=reuse) variable_scope_context_manager = getattr(obj, "_enter_variable_scope", default_context_manager) graph = tf.get_default_graph() if graph not in initialized_variable_scopes: initialized_variable_scopes[graph] = set() initialized_variable_scopes_for_graph = initialized_variable_scopes[graph] # Temporarily enter the variable scope to capture it with variable_scope_context_manager() as tmp_variable_scope: variable_scope = tmp_variable_scope reuse = variable_scope.name in initialized_variable_scopes_for_graph # Enter the pure variable scope with reuse correctly set with variable_scope_ops._pure_variable_scope( # pylint:disable=protected-access variable_scope, reuse=reuse) as pure_variable_scope: # Force tf.name_scope to treat variable_scope.original_name_scope as # an "absolute" scope name so we can re-enter it. name_scope = variable_scope.original_name_scope if name_scope[-1] != "/": name_scope += "/" with tf.name_scope(name_scope): sub_scope = to_snake_case(method.__name__) with tf.name_scope(sub_scope) as scope: if hasattr(obj, "_capture_variables"): with obj._capture_variables(): # pylint: disable=protected-access out_ops = method(obj, *args, **kwargs) else: out_ops = method(obj, *args, **kwargs) initialized_variable_scopes_for_graph.add(pure_variable_scope.name) try: # If `obj` is a Sonnet module, let it know it's been connected # to the TF graph method_positional_args = [obj] + list(args) obj._add_connected_subgraph( # pylint: disable=protected-access method, out_ops, scope, *method_positional_args, **kwargs) except AttributeError: pass return out_ops
def call_method(method, obj, args, kwargs): """Calls `method` with a variable scope whose reuse flag is set correctly. The first time the wrapper is called it creates a `(tf.Graph, tf.VariableScope)` key and checks it for membership in `initialized_variable_scopes`. The check is `False` if and only if this is the first time the wrapper has been called with the key, otherwise the check is `True`. The result of this check is used as the `reuse` flag for entering the provided variable scope before calling `method`. Here are two examples of how to use the reuse_variables decorator. 1. Decorate an arbitrary instance method with a `variable_scope` attribute: ```python class Reusable(object): def __init__(self, name): with tf.variable_scope(None, default_name=name) as vs: self.variable_scope = vs @snt.reuse_variables def add_a(self, input_tensor): a = tf.get_variable("a", shape=input_tensor.get_shape()) return a + input_tensor obj = Reusable("reusable") x = tf.constant(5.0) out1 = obj.add_a(x) out2 = obj.add_a(x) # out1 == out2 ``` 2. Decorating a snt.AbstractModule instance method: ```python class ReusableModule(snt.AbstractModule): @snt.reuse_variables def add_a(self, input_tensor): a = tf.get_variable("a", shape=input_tensor.get_shape()) return a + input_tensor # We don't need @snt.reuse_variables here because build is wrapped by # `tf.make_template` inside `snt.AbstractModule`. def _build(self, input_tensor): b = tf.get_variable("b", shape=input_tensor.get_shape()) return b + self.add_a(input_tensor) obj = Reusable("reusable") x = tf.constant(5.0) out1 = obj(x) out2 = obj(x) # out1 == out2 ``` Args: method: The method to wrap. obj: The object instance passed to the wrapped method. args: The positional arguments (Tensors) passed to the wrapped method. kwargs: The keyword arguments passed to the wrapped method. Returns: Output of the wrapped method. Raises: ValueError: If no variable scope is provided or if `method` is a method and a variable_scope keyword argument is also provided. """ # If @reuse_variables is combined with @property, obj is passed in args # and method is still unbound at this stage. if obj is None: obj = args[0] def default_context_manager(reuse=None): variable_scope = obj.variable_scope return tf.variable_scope(variable_scope, reuse=reuse) variable_scope_context_manager = getattr(obj, "_enter_variable_scope", default_context_manager) with ops.init_scope(): # We need `init_scope` incase we're running inside a defun. In that case # what we want is information about where the function will be called not # where the function is being built. graph = tf.get_default_graph() will_call_in_eager_context = tf.executing_eagerly() if will_call_in_eager_context: initialized_variable_scopes = initialized_variable_scopes_eager else: if graph not in initialized_variable_scopes_graph: initialized_variable_scopes_graph[graph] = set() initialized_variable_scopes = initialized_variable_scopes_graph[graph] # Temporarily enter the variable scope to capture it with variable_scope_context_manager() as tmp_variable_scope: variable_scope = tmp_variable_scope reuse = variable_scope.name in initialized_variable_scopes # Enter the pure variable scope with reuse correctly set with variable_scope_ops._pure_variable_scope( # pylint:disable=protected-access variable_scope, reuse=reuse) as pure_variable_scope: current_name_scope = tf.get_default_graph().get_name_scope() # Force tf.name_scope to treat current_name_scope as an "absolute" scope # so we can re-enter it. if current_name_scope and current_name_scope[-1] != "/": current_name_scope += "/" with tf.name_scope(current_name_scope): module_name = pure_variable_scope.name method_name = to_snake_case(method.__name__) method_name_scope = "{}/{}".format(module_name, method_name) with tf.name_scope(method_name_scope) as scope: if hasattr(obj, "_capture_variables"): with obj._capture_variables(): # pylint: disable=protected-access out_ops = method(*args, **kwargs) else: out_ops = method(*args, **kwargs) initialized_variable_scopes.add(pure_variable_scope.name) try: # If `obj` is a Sonnet module, let it know it's been connected # to the TF graph. obj._is_connected = True # pylint: disable=protected-access if not tf.executing_eagerly(): obj._add_connected_subgraph( # pylint: disable=protected-access method, out_ops, scope, *args, **kwargs) except AttributeError: pass return out_ops