Module API

The Module class is a fundamental building block in MLLM’s neural network framework. It serves as a container for neural network components and provides functionalities for parameter management, device placement, and execution.

#include "mllm/nn/Module.hpp"

Base Class

class Module

Base class for neural network modules. Modules can contain other modules or layers and manage their parameters and execution.

Constructors

Module::Module()

Default constructor.

explicit Module::Module(const ModuleImpl::ptr_t &impl)

Constructor with ModuleImpl pointer.

Parameters:

impl – Shared pointer to ModuleImpl instance

explicit Module::Module(const std::string &name)

Constructor with module name.

Parameters:

name – Name of the module

Core Methods

ModuleImpl::ptr_t Module::impl() const

Get the underlying ModuleImpl pointer.

Returns:

Shared pointer to ModuleImpl

void Module::to(DeviceTypes device_type)

Move the module and its parameters to specified device.

Parameters:

device_type – Target device type (kCPU, kCUDA, etc.)

template<typename T, typename ...Args>
auto Module::reg(const std::string &name, Args&&... args)

Register a sub-module or layer to this module.

Parameters:

  • name – Name of the sub-module or layer

  • args – Arguments for constructing the sub-module or layer

Returns:

Registered sub-module or layer

template<typename ...Args>
std::vector<Tensor> Module::operator()(Args&&... args)

Execute the module with given inputs.

Parameters:

args – Input tensors and other arguments

Returns:

Output tensors

void Module::load(const ParameterFile::ptr_t &param_file)

Load parameters from a parameter file.

Parameters:

param_file – Shared pointer to ParameterFile

virtual std::vector<Tensor> Module::forward(const std::vector<Tensor> &inputs, const std::vector<AnyValue> &args)

Forward pass of the module. Should be implemented by derived classes.

Parameters:

  • inputs – Input tensors

  • args – Additional arguments

Returns:

Output tensors

Utility Methods

void Module::__fmt_print(std::stringstream &ss) const

Format print information about the module.

Parameters:

ss – String stream to write formatted output

std::vector<Tensor> Module::__main(const std::vector<Tensor> &inputs, const std::vector<AnyValue> &args)

Main execution method that handles preprocessing and postprocessing.

Parameters:

  • inputs – Input tensors

  • args – Additional arguments

Returns:

Output tensors

void Module::__send_graph_begin(const std::vector<Tensor> &inputs)

Send graph begin signal for execution tracing.

Parameters:

inputs – Input tensors

void Module::__send_graph_end(const std::vector<Tensor> &inputs)

Send graph end signal for execution tracing.

Parameters:

inputs – Input tensors

std::vector<Tensor> Module::__trace(const std::vector<Tensor> &inputs, const std::vector<AnyValue> &args)

Trace execution for compilation or analysis.

Parameters:

  • inputs – Input tensors

  • args – Additional arguments

Returns:

Output tensors

ParameterFile::ptr_t Module::params(ModelFileVersion v)

Get parameters of the module.

Parameters:

v – Model file version

Returns:

Shared pointer to ParameterFile

void Module::registerBuffer(const std::string &name, const Tensor &tensor)

Register a buffer tensor that won’t be saved with parameters.

Parameters:

  • name – Name of the buffer

  • tensor – Buffer tensor

Tensor Module::getBuffer(const std::string &name)

Get a registered buffer tensor.

Parameters:

name – Name of the buffer

Returns:

Buffer tensor

std::string Module::getModuleName() const

Get the full name of the module.

Returns:

Module name