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,
@@ -156,7 +157,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 +177,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 +298,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 +444,33 @@ def found_dup_callback(x, kc):
             vs = vs.cont_prune_key_chain(dup_kc)
         return vs, keychain_mappings
 
+    def _set_buffers(self, buffers):
+        """
+        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])
+                else:
+                    setattr(self, buffer, buffers[buffer])
+            else:
+                if hasattr(self, "buffers"):
+                    self.buffers.update({buffer: buffers[buffer]})
+                else:
+                    setattr(self, "buffers", {buffer: buffers[buffer]})
+                setattr(self, buffer, buffers[buffer])
+
     # Overridable #
 
     # noinspection PyMethodMayBeStatic,PyUnusedLocal
@@ -494,7 +537,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 +559,10 @@ def _call(self, *args, v=None, **kwargs):
                 from_call=True,
                 dtype=_get_first_array(*args, **kwargs).dtype,
             )
+        if buffers:
+            buffers_orig = self.buffers.copy()
+            self.buffers = {}
+            self._set_buffers(buffers)
         if v is not None:
             v_orig = self.v
             self.v = (
@@ -525,7 +572,11 @@ def _call(self, *args, v=None, **kwargs):
             )
             ret = self._forward_with_tracking(*args, **kwargs)
             self.v = v_orig
+            if buffers:
+                self.buffers = {}
+                self._set_buffers(buffers_orig)
             return ret
+
         elif hasattr(self.__call__, "wrapped"):
             return self.__call__(*args, **kwargs)
         return self._forward_with_tracking(*args, **kwargs)
@@ -536,6 +587,7 @@ def __call__(
         self,
         *args,
         v=None,
+        buffers=None,
         stateful=None,
         arg_stateful_idxs=None,
         kwarg_stateful_idxs=None,
@@ -597,7 +649,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 +676,7 @@ def build(
         device=None,
         dtype=None,
         dynamic_backend=None,
+        buffers=None,
         **kwargs,
     ):
         """
@@ -647,7 +700,6 @@ def build(
         """
         self._dev = ivy.default(device, self._dev)
         # 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,8 +801,18 @@ 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 register_buffer(self, var_name, value):
+        """Set the buffer at any place within the class."""
+        self._set_buffers({var_name: value})
+
     def __repr__(self):
         return object.__repr__(self)
 
@@ -798,9 +860,26 @@ def __getattribute__(self, name):
                 self._build_and_return_v(
                     self._args, dynamic_backend=self._dynamic_backend, **self._kwargs
                 )
-
+        if name != "buffers":
+            if hasattr(self, "buffers"):
+                if name in self.buffers:
+                    return self.buffers[name]
         return super().__getattribute__(name)
 
+    def __setattr__(self, name, value):
+        if hasattr(self, "buffers"):
+            if name in self.buffers:
+                self.buffers[name] = value
+                return
+        return super().__setattr__(name, value)
+
+    def __delattr__(self, name):
+        if hasattr(self, "buffers"):
+            if name in self.buffers:
+                del self.buffers[name]
+        else:
+            super().__delattr__(name)
+
     def compile(
         self,
         args: Optional[Tuple] = None,

ivy_tests/test_ivy/test_stateful/test_modules.py

@@ -1059,3 +1059,33 @@ def loss_fn(v_):
     assert ivy.Container.all(loaded_module.v == module.v).cont_all_true()
 
     os.remove(save_filepath)
+
+
+class ModuleWithBuffer(ivy.Module):
+    def __init__(self, *args, **kwargs):
+        pass
+
+    def _forward(*args, **kwargs):
+        pass
+
+
+@given(
+    buffer=st.just(
+        [
+            {
+                "var1": [
+                    ivy.ones((1, 2)),
+                ]
+            }
+        ]
+    )
+)
+def test_get_buffers(buffer):
+    module = ModuleWithBuffer()
+    buffers = {}
+    for item in buffer:
+        buffers.update(item)
+        for key in item:
+            module.register_buffer(key, item[key])
+
+    assert module.buffers == buffers