Module
- class ivy.stateful.module.Module(*args, device=None, v=None, build_mode='on_init', compile_on_next_step=False, store_vars=True, stateful=None, arg_stateful_idxs=None, kwarg_stateful_idxs=None, fallback_to_non_compiled=False, with_partial_v=False, devices=None, dtype=None, **kwargs)[source]
Bases:
ModuleConverters
,ModuleHelpers
Module is a base class for deriving trainable modules.
- __call__(*args, v=None, with_grads=None, stateful=None, arg_stateful_idxs=None, kwarg_stateful_idxs=None, track_submod_rets=False, submod_depth=None, submods_to_track=None, track_submod_call_order=False, expected_submod_rets=None, **kwargs)[source]
Forward an input through current module.
- Parameters
v – If given, use this container as internal varibles temporarily. Default is
None
.with_grads – If True, forward this pass with gradients.
track_submod_rets – If True, will track the returns of submodules.
submod_depth – The depth of tracked submodules.
submods_to_track – If given, will only track submodules in submods_to_track.
track_submod_call_order – If True, will track the call order of submodules.
expected_submod_rets – If given, will raise exception if submodule returns are different from expected returns.
- Returns
ret
- __init__(*args, device=None, v=None, build_mode='on_init', compile_on_next_step=False, store_vars=True, stateful=None, arg_stateful_idxs=None, kwarg_stateful_idxs=None, fallback_to_non_compiled=False, with_partial_v=False, devices=None, dtype=None, **kwargs)[source]
Initialize Ivy layer, which is a stateful object consisting of trainable variables.
- Parameters
device – device on which to create the module’s variables ‘cuda:0’, ‘cuda:1’, ‘cpu’ etc. (Default value = None)
v – Ivy container of trainable variables. Created internally by default.
build_mode – How the Module is built, either on initialization (now), explicitly by the user by calling build(), or the first time the __call__ method is run. Default is on initialization.
compile_on_next_step – Whether to compile the network on the next forward pass. Default is
False
.store_vars – Whether or not to store the variables created. Default is
True
.stateful – The constant id stateful items to track as part of the forward pass. Used when graph compiling, default is
None
.arg_stateful_idxs – The nested argument indices of stateful items to track as part of the forward pass. Used when graph compiling, default is
None
.kwarg_stateful_idxs – The nested keyword argument indices of stateful items to track as part of the forward pass. Used when graph compiling, default is
None
.fallback_to_non_compiled – Whether to fall back to non-compiled forward call in the case that an error is raised during the compiled forward pass. Default is
True
.with_partial_v – Whether to allow partial specification of variables. Default is
False
.devices – devices on which to distribute the module’s variables ‘cuda:0’, ‘cuda:1’, ‘cpu’ etc. (Default value = None)
- build(*args, from_call=False, device=None, dtype=None, **kwargs)[source]
Build the internal layers and variables for this module.
- Parameters
from_call – If True, denote that this build is triggered by calling. Otherwise, triggered by initializing the module. Default is
False
.device – The device we want to build module on. None for default device. Default is
None
.dtype – The data type for building the module. Default is
None
.
- Returns
ret – True for successfully built a module.
- property build_mode
- property built_