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_
save_weights(weights_path, /)[source]

Save the weights on the Module.

Parameters

weights_path – The hdf5 file for saving the weights.

Returns

None


Supported Frameworks: