Module#

Base class for deriving trainable modules.

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, dynamic_backend=None, **kwargs)[source]#

Bases: ModuleConverters, ModuleHelpers

Module is a base class for deriving trainable modules.

__call__(*args, v=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.

  • 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, dynamic_backend=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, dynamic_backend=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_#
compile(args=None, kwargs=None, **compile_kwargs)[source]#

Compile the ivy.Module’s _unified_ivy_graph or _call method to the target backend.

Parameters:
  • args (Optional[Tuple]) – arguments used to compile. Defaults to None. (default: None)

  • kwargs (Optional[Dict]) – keyword arguments used to compile. Defaults to None. (default: None)

  • compile_kwargs – keyword arguments passed to the compile function.

save_weights(weights_path, /)[source]#

Save the weights on the Module.

Parameters:

weights_path – The hdf5 file for saving the weights.

Returns:

None

show_graph(randomness_factor=0.1, save_to_disk=False, notebook=False, with_edge_labels=True, with_arg_labels=True, with_output_labels=True, output_connected_only=True, highlight_subgraph=None, fname=None)[source]#

This should have hopefully given you an overview of the module submodule, if you have any questions, please feel free to reach out on our discord in the module channel or in the module forum!