Stateful changes 2 #21246
Stateful changes 2 #21246
Conversation
…ation on demand, modified container methods container identity methods to have build_callable and assert and assign arguments to build any callables as modules and assign them immediately
…oic/ivy into stateful-changes-2
…oic/ivy into stateful-changes-2
…s, rather than having it in the main module class
…d some corner cases in Module _init_var to account for absence and deleting it when done
…,_get_buffers,_register_buffers,_buffer_tracker, _create_buffers added. Alos _call was modified to account for the 'buffers' passed during call
|
Thanks for contributing to Ivy! 😊👏 |
ivy/stateful/module.py
Outdated
| else: | ||
| setattr(self, "buffers", {buffer}) | ||
|
|
||
| def _register_buffers(self, var_name, value): |
There was a problem hiding this comment.
make this public, and remove variables.
There was a problem hiding this comment.
could you please rename register_buffers to register_buffer, set_buffers to set_buffer and others and make the necessary ones public?
There was a problem hiding this comment.
Hey @RickSanchezStoic, just made a few suggestions, seems like we need to clean it up a bit. Thanks 😄
ivy/stateful/module.py
Outdated
| @@ -144,6 +145,9 @@ def __init__( | |||
| self._target = None | |||
| self._lazy_compiled = False | |||
| self._dynamic_backend = dynamic_backend | |||
| if hasattr(self, "_create_buffers"): | |||
There was a problem hiding this comment.
Seems like we're still using _create_buffers and _buffer_tracker, could you please remove it as we'd only need to allow users to register buffers "on-spot"?
ivy/stateful/module.py
Outdated
|
|
||
| return | ||
| self.build(*args, dynamic_backend=dynamic_backend, **kwargs) | ||
| self.build(*args, dynamic_backend=dynamic_backend, buffers=buffers**kwargs) |
There was a problem hiding this comment.
syntax error, , to be added
There was a problem hiding this comment.
this throws an error right?
ivy/stateful/module.py
Outdated
| wrapper.buffers_tracked = True | ||
| return wrapper | ||
|
|
||
| def _get_buffers(self): |
There was a problem hiding this comment.
I'm not sure why we need this method actually, given that we are already using _find_buffers while building, won't self.buffers be a container or arrays anyway?
ivy/stateful/module.py
Outdated
| else: | ||
| setattr(self, "buffers", {buffer}) | ||
|
|
||
| def _register_buffers(self, var_name, value): |
There was a problem hiding this comment.
could you please rename register_buffers to register_buffer, set_buffers to set_buffer and others and make the necessary ones public?
…oic/ivy into stateful-changes-2
…l updates, fetches . deletions performed on the self.buffers, updated tests
There was a problem hiding this comment.
Hey @RickSanchezStoic, just requested a final few changes, thanks 😄
ivy/stateful/module.py
Outdated
|
|
||
| return | ||
| self.build(*args, dynamic_backend=dynamic_backend, **kwargs) | ||
| self.build(*args, dynamic_backend=dynamic_backend, buffers=buffers**kwargs) |
ivy/stateful/module.py
Outdated
| @@ -646,8 +699,9 @@ def build( | |||
| True for successfully built a module. | |||
| """ | |||
| self._dev = ivy.default(device, self._dev) | |||
| # build buffers if any | |||
| self._create_buffers() | |||
There was a problem hiding this comment.
Given that ivy.Module._create_buffers has been removed, we should also remove this
ivy/stateful/module.py
Outdated
| @@ -798,9 +862,27 @@ def __getattribute__(self, name): | |||
| self._build_and_return_v( | |||
| self._args, dynamic_backend=self._dynamic_backend, **self._kwargs | |||
| ) | |||
|
|
|||
| if name == "buffers": | |||
There was a problem hiding this comment.
This if block is redundant right?
ivy/stateful/module.py
Outdated
| if name in self.buffers: | ||
| self.buffers[name] = value | ||
| return | ||
|
|
There was a problem hiding this comment.
Final couple of changes left from the previous review, thanks @RickSanchezStoic 😄
ivy/stateful/module.py
Outdated
|
|
||
| return | ||
| self.build(*args, dynamic_backend=dynamic_backend, **kwargs) | ||
| self.build(*args, dynamic_backend=dynamic_backend, buffers=buffers**kwargs) |
There was a problem hiding this comment.
this throws an error right?
Co-authored-by: Rishabh Kumar <[email protected]> Adds buffer support to Module
Added buffer variables support for ivy Module. Methods such as _set_buffers,_get_buffers,_register_buffers,_buffer_tracker, _create_buffers added. Also _call was modified to account for the 'buffers' passed during the call.