Stateful changes 2

ivy/data_classes/container/base.py

@@ -2845,7 +2845,7 @@ def cont_overwrite_at_key_chains(
                 )
             else:
                 return_dict[k] = v
-        return ivy.Container(return_dict, **self._config)
+        return return_dict
 
     def cont_prune_keys(self, query_keys, ignore_none=True):
         """

ivy/stateful/module.py

@@ -49,6 +49,7 @@ def __init__(
         *args,
         device=None,
         v=None,
+        buffers=None,
         build_mode="on_init",
         compile_on_next_step=False,
         store_vars=True,
@@ -144,6 +145,9 @@ def __init__(
         self._target = None
         self._lazy_compiled = False
         self._dynamic_backend = dynamic_backend
+        if hasattr(self, "_create_buffers"):
+            self._create_buffers = self._buffer_tracker(self._create_buffers)
+
         if build_mode != "on_init":
             return
         if hasattr(Module, "_init_var"):
@@ -156,7 +160,12 @@ def __init__(
 
                 if v or with_partial_v:
                     # build only if `v` or `with_partial_v`
-                    self.build(*args, dynamic_backend=dynamic_backend, **kwargs)
+                    self.build(
+                        *args,
+                        dynamic_backend=dynamic_backend,
+                        buffers=buffers,
+                        **kwargs,
+                    )
                 # we don't want to delete the class variable now
                 # since there could be other child modules
                 return
@@ -171,16 +180,19 @@ def __init__(
                 # move on
                 Module._init_var.pop()
                 return
-            self.build(*args, dynamic_backend=dynamic_backend, **kwargs)
+            self.build(
+                *args, dynamic_backend=dynamic_backend, buffers=buffers, **kwargs
+            )
             if Module._init_var[-1] == self.__class__.__name__:
                 # you delete it, only if this is the class that caused it's creation
                 Module._init_var.pop()
 
             # do a final check if _init_var  becomes empty, then delete it all together
-            del Module._init_var
+            if not Module._init_var:
+                del Module._init_var
 
             return
-        self.build(*args, dynamic_backend=dynamic_backend, **kwargs)
+        self.build(*args, dynamic_backend=dynamic_backend, buffers=buffers**kwargs)
 
     # Private #
     # --------#
@@ -289,6 +301,13 @@ def _build_and_return_v(self, *args, **kwargs):
     def _find_child_objects(self, /, *, obj=None, _visited=None):
         pass
 
+    def _find_buffers(self):
+        for obj in self.__dict__.keys():
+            if isinstance(getattr(self, obj), ivy.Module):
+                # simply fetch it's buffer
+                if hasattr(getattr(self, obj), "buffers"):
+                    self.buffers.update({obj: getattr(self, obj).buffers})
+
     @staticmethod
     def _extract_v(v, keychain_mappings: dict, orig_key_chain, /):
         """
@@ -428,6 +447,76 @@ def found_dup_callback(x, kc):
             vs = vs.cont_prune_key_chain(dup_kc)
         return vs, keychain_mappings
 
+    def _buffer_tracker(self, func):
+        """Tracks the variables defined as buffer variables and stores them."""
+
+        def wrapper(*args, **kwargs):
+            initial_object_snapshot = self.__dict__.copy()
+            # initialise the buffers
+            func(*args, **kwargs)
+            final_object_snapshot = self.__dict__
+            getattr(self, "buffers", {}).update(
+                set(final_object_snapshot.keys()).difference(
+                    set(initial_object_snapshot.keys())
+                )
+            )
+
+        wrapper.buffers_tracked = True
+        return wrapper
+
+    def _get_buffers(self):
+        """Return the buffer variables, if any, of the given Module class instance."""
+        if self.buffers:
+            buffer_dict = {}
+            for buffer in self.buffers:
+                buffer_dict[buffer] = (
+                    getattr(self, buffer)
+                    if not isinstance(getattr(self, buffer), ivy.Module)
+                    else getattr(self, buffer)._get_buffers()
+                )
+            return buffer_dict
+        return {}
+
+    def _set_buffers(self, buffers, override=False, nesting=None):
+        """
+        Set the buffers of the given class instance, according to the buffers passed.
+
+        Parameters
+        ----------
+        buffers
+            a dictionary with variable names and corresponding values
+
+        override
+            if true, sets the variable as an attribute even if it doesn't exist
+        """
+        for buffer in buffers:
+            if hasattr(self, buffer):
+                # check if this value is another nested dictionary, if yes
+                # we recurse
+                if isinstance(buffers[buffer], dict):
+                    getattr(self, buffer)._set_buffers(
+                        buffers=buffers[buffer], override=override
+                    )
+                else:
+                    setattr(self, buffer, buffers[buffer])
+            elif not override or (override and isinstance(buffers[buffer], dict)):
+                # do a quick check to see if this is nested case
+                # we can't set buffers for module objects not yet
+                # created, even if override=True
+                raise ivy.exceptions.IvyNotImplementedException(
+                    f"{buffer} hasn't been defined for the given Module structure"
+                )
+            else:
+                setattr(self, buffer, buffers[buffer])
+                if hasattr(self, "buffers"):
+                    self.buffers.update({buffer})
+                else:
+                    setattr(self, "buffers", {buffer})
+
+    def _register_buffers(self, var_name, value):
+        """Set the buffer variables at any place within the class."""
+        self._set_buffers({var_name: value}, override=True)
+
     # Overridable #
 
     # noinspection PyMethodMayBeStatic,PyUnusedLocal
@@ -448,6 +537,17 @@ def _create_variables(self, *, device=None, dtype=None):
         """
         return {}
 
+    def _create_buffers(self):
+        """
+        Create buffers for this class.
+
+        Returns
+        -------
+        ret
+            An empty set.
+        """
+        return {}
+
     def _build(self, *args, **kwargs) -> bool:
         """
         Build the internal layers and variables for this module. Overridable.
@@ -494,7 +594,7 @@ def _forward_with_tracking(self, *args, **kwargs):
             self._check_submod_ret()
         return ret
 
-    def _call(self, *args, v=None, **kwargs):
+    def _call(self, *args, v=None, buffers=None, **kwargs):
         """
         Compute forward pass of the layer, treating layer instance as callable function.
 
@@ -516,6 +616,9 @@ def _call(self, *args, v=None, **kwargs):
                 from_call=True,
                 dtype=_get_first_array(*args, **kwargs).dtype,
             )
+        if buffers:
+            buffers_orig = self._get_buffers()
+            self._set_buffers(buffers)
         if v is not None:
             v_orig = self.v
             self.v = (
@@ -525,6 +628,8 @@ def _call(self, *args, v=None, **kwargs):
             )
             ret = self._forward_with_tracking(*args, **kwargs)
             self.v = v_orig
+            if buffers:
+                self._set_buffers(buffers_orig)
             return ret
         elif hasattr(self.__call__, "wrapped"):
             return self.__call__(*args, **kwargs)
@@ -536,6 +641,7 @@ def __call__(
         self,
         *args,
         v=None,
+        buffers=None,
         stateful=None,
         arg_stateful_idxs=None,
         kwarg_stateful_idxs=None,
@@ -597,7 +703,7 @@ def __call__(
 
         # convert variables to native arrays so that they can be tracked
         v = ivy.to_native(v)
-        ret = self._call(*args, v=v, **kwargs)
+        ret = self._call(*args, v=v, buffers=buffers, **kwargs)
         self._unset_submod_flags()
         return ret
 
@@ -624,6 +730,7 @@ def build(
         device=None,
         dtype=None,
         dynamic_backend=None,
+        buffers=None,
         **kwargs,
     ):
         """
@@ -646,8 +753,9 @@ def build(
             True for successfully built a module.
         """
         self._dev = ivy.default(device, self._dev)
+        # build buffers if any
+        self._create_buffers()
         # return False if not from_call but build_mode is on_call
-
         if not from_call and self._build_mode == "on_call":
             return self.v
         if dtype:
@@ -749,6 +857,12 @@ def build(
         if not self._store_vars:
             # ToDo: verify variables in self.v are released once this method exits
             self.v = ivy.Container()
+
+        # once all variables built, find and assign buffers
+        if buffers:
+            self._set_buffers(buffers=buffers)
+            self._find_buffers()
+
         return v_ret if bool(v_ret) or isinstance(built, bool) else built
 
     def __repr__(self):