A component that contains and manages child components. More...

Inheritance diagram for Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >:

Public Types
using	ComponentBase = Component<TDeviceType, TPrecision>
using	ComponentPtr = std::shared_ptr<Component<TDeviceType, TPrecision>>

Public Member Functions
	CompositeComponent (CompositeComponent &&) noexcept=default
	CompositeComponent (const CompositeComponent &)=delete
	CompositeComponent (const std::string &name)
	Construct composite component with name.
CompositeComponent &	addComponent (ComponentPtr component)
	Add a pre-constructed child component (chainable).
size_t	childCount () const noexcept
	Get the number of direct children.
void	clearComponents ()
	Clear all child components.
ComponentPtr	findComponent (const std::string &path) const
	Resolve a dot-separated component path within this composite.
ComponentPtr	getComponent (const std::string &name) const
	Retrieve a direct child component by name.
const std::vector< ComponentPtr > &	getComponents () const
	Get all child components in insertion order.
DeviceId	getDeviceId () const override
	Get the compute device for this composite.
std::vector< ITensor * >	getGradients () const override
	Get all parameter gradients from all children.
std::vector< ITensor * >	getParameters () const override
	Get all parameters from all children.
bool	hasChildren () const noexcept
	Check if this composite has any children.
bool	hasComponent (const std::string &name) const
	Check if a named child component exists.
CompositeComponent &	operator= (CompositeComponent &&) noexcept=default
CompositeComponent &	operator= (const CompositeComponent &)=delete
dim_t	parameterCount () const override
	Count parameters across all children.
bool	removeComponent (const std::string &name)
void	saveFlatTensors (Serialization::SafeTensorsWriter &writer, const std::string &prefix, Serialization::TensorSavePass pass) const override
	Recurse into children, extending the flat dotted prefix.
void	synchronize () override
	Synchronize all child components.
std::string	toString () const override
	Generate a human-readable description.
ComponentPtr	tryFindComponent (const std::string &path) const
	Try to resolve a dot-separated component path within this composite.
Public Member Functions inherited from Mila::Dnn::Component< TDeviceType, TPrecision >
	Component (const std::string &name)
	Construct component with required name identifier.
virtual void	build (const BuildContext &context) final
	Build the component with the provided BuildContext (canonical overload).
virtual MemoryStats	getMemoryStats () const =0
	Return the current memory allocation breakdown for this component.
const std::string	getName () const
	Get the component's name identifier.
virtual std::vector< std::string >	getParameterNames () const
	List all available parameter names for this component.
virtual MemoryStats	getRequiredMemory (const BuildContext &context) const
	Report what build( context ) would allocate, without allocating it.
TrainingMode	getTrainingMode () const noexcept
	The current runtime behavioral mode of this Component.
virtual const ComponentType	getType () const =0
	Get the component type identifier.
virtual bool	isBuilt () const final
	Returns true if build() has completed successfully.
virtual void	loadParameter (const std::string &, const Serialization::ITensorBlob &)
	Load a parameter from serialized tensor data.
void	setTrainingMode (TrainingMode mode)
	Set the runtime behavioral mode for this Component.
virtual void	zeroGradients ()
	Clear all model-owned gradients for this component.

Protected Member Functions
template<typename TComponent>
std::shared_ptr< TComponent >	getComponentAs (const std::string &name) const
	Retrieve a typed child component by name.
void	load_ (ModelArchive &archive, SerializationMode mode) override
	Restore children from their nested scopes, mirroring save_().
void	onExecutionContextSet () override
	Hook invoked after ExecutionContext is set.
void	onTrainingModeChanging (TrainingMode training_mode) override
	Hook invoked when training mode is about to change.
virtual void	optimize ()
	Virtual hook for graph optimization after construction.
void	requireSerializableParameters () const override
	No-op override: a composite names no parameters of its own.
void	save_ (ModelArchive &archive, SerializationMode mode) const override
Protected Member Functions inherited from Mila::Dnn::Component< TDeviceType, TPrecision >
IExecutionContext *	getExecutionContext () const
	Get the shared execution context.
bool	hasExecutionContext () const noexcept
	Check if execution context has been set.
template<TensorDataType TParameterPrecision, typename TMemoryResource>
void	loadParameterFromBlob (const std::string &param_name, const Serialization::ITensorBlob &blob, Tensor< TParameterPrecision, TMemoryResource > &target, const shape_t &expected_shape)
	Load a tensor blob into a parameter tensor with validation.
virtual void	onBuilding (const BuildContext &)
	Hook invoked by build() to allocate component buffers.
template<TensorDataType TParameterPrecision, typename TMemoryResource>
void	saveParameterToArchive (ModelArchive &archive, const std::string &parameter_name, const Tensor< TParameterPrecision, TMemoryResource > &parameter) const
	Write one parameter tensor into the archive under "tensors/<name>".
template<TensorDataType TParameterPrecision, typename TMemoryResource>
void	saveParameterToWriter (Serialization::SafeTensorsWriter &writer, const std::string &flat_name, const Tensor< TParameterPrecision, TMemoryResource > &parameter, Serialization::TensorSavePass pass) const
	Drive one parameter through one pass of a flat safetensors save.
void	setExecutionContext (IExecutionContext *context)
	Set the execution context for this component.

Friends
class	ComponentFactory

Additional Inherited Members
Static Public Member Functions inherited from Mila::Dnn::Component< TDeviceType, TPrecision >
static constexpr DeviceType	getDeviceType ()
	Compile-time device type for this component instance.
static constexpr TensorDataType	getPrecision () noexcept
	Compile-time tensor precision for this component instance.
Protected Types inherited from Mila::Dnn::Component< TDeviceType, TPrecision >
using	HostStagingMemoryResource
	Host memory a device-resident parameter stages through.
Protected Attributes inherited from Mila::Dnn::Component< TDeviceType, TPrecision >
BuildContext	build_context_ { shape_t{ 1 }, RuntimeMode::Training }
	The BuildContext stored at build time.

Detailed Description

template<DeviceType TDeviceType, TensorDataType TPrecision>
class Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >

A component that contains and manages child components.

CompositeComponent is a device-parameterized abstract container that manages child component lifecycle, aggregates operations (parameters, gradients, training mode), and provides context propagation. Derived types implement execution semantics (forward/backward) and architecture definition (createGraph()).

Architecture Philosophy:

Context-independent graph creation: Architecture defined without device knowledge
Three-phase lifecycle: Graph creation -> Context binding -> Shape binding
Automatic context propagation: Base class propagates context to all children
Component-owns-name: Children manage their own identity via getName()

NOTE:

getComponent() performs direct child lookup only.
For path-based / full-graph resolution use findComponent() or tryFindComponent() on the Network (root) or call findComponent() on a composite to resolve dot-separated paths (e.g. "encoder.mlp.fc1").

Constructor & Destructor Documentation

◆ CompositeComponent()

template<DeviceType TDeviceType, TensorDataType TPrecision>

Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::CompositeComponent ( const std::string & name )

inlineexplicit

Construct composite component with name.

The composite component is named for identification in hierarchical structures. Derived classes should call createGraph() from their constructor to define the architecture graph (context-independent).

All child components added via addComponent() will receive ExecutionContext automatically when the composite receives its context (via onExecutionContextSet).

Parameters

name	Component name identifier (mandatory)

Exceptions

std::invalid_argument if name is not a valid identifier

Member Function Documentation

◆ addComponent()

template<DeviceType TDeviceType, TensorDataType TPrecision>

CompositeComponent & Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::addComponent ( ComponentPtr component )

inline

Add a pre-constructed child component (chainable).

Registers a component that was constructed externally (typically by the derived class in its createGraph() method). The component's getName() is used as the lookup key.

Components are expected to be created in shared mode (no ExecutionContext). Context will be automatically propagated to all children when this composite receives its context via onExecutionContextSet().

Usage pattern in derived class:

void MLP::createGraph()
{
    auto fc1 = std::make_shared<LinearType>(config, std::nullopt);
    fc1->setName(this->getName() + ".fc1");
    this->addComponent(fc1);
    // ... more components
}

Parameters

component Shared pointer to the constructed component

Returns: Reference to *this for method chaining

Exceptions

std::runtime_error	if called after build()
std::invalid_argument	if component is null
std::invalid_argument	if component name already exists
std::invalid_argument	if component already has its own ExecutionContext

◆ childCount()

template<DeviceType TDeviceType, TensorDataType TPrecision>

size_t Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::childCount ( ) const

inlinenoexcept

Get the number of direct children.

Returns: Number of child components

◆ clearComponents()

template<DeviceType TDeviceType, TensorDataType TPrecision>

void Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::clearComponents ( )

inline

Clear all child components.

Exceptions

std::runtime_error if called after build()

◆ findComponent()

template<DeviceType TDeviceType, TensorDataType TPrecision>

ComponentPtr Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::findComponent ( const std::string & path ) const

inline

Resolve a dot-separated component path within this composite.

Supports both relative paths ("lenc.wte") and absolute paths ("gpt2.lenc.wte"). If path starts with this component's name, strips it before searching.

◆ getComponent()

template<DeviceType TDeviceType, TensorDataType TPrecision>

ComponentPtr Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::getComponent ( const std::string & name ) const

inline

Retrieve a direct child component by name.

This performs direct (non-recursive) lookup of immediate children only. Use findComponent() to resolve dot-separated paths across the subgraph.

Parameters

name	Name of the direct child component

Returns: Shared pointer to the component

Exceptions

std::out_of_range if the direct child is not found

◆ getComponentAs()

template<DeviceType TDeviceType, TensorDataType TPrecision>

template<typename TComponent>

std::shared_ptr< TComponent > Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::getComponentAs ( const std::string & name ) const

inlineprotected

Retrieve a typed child component by name.

Helper method for derived composites (like MLP) that need to cache typed pointers to children in their onBuilding() hook. Performs dynamic_pointer_cast and validates the cast succeeded.

Note: This resolves direct children only. For full-path resolution use findComponent() on the appropriate root composite or Network.

Template Parameters

TComponent Expected component type

Parameters

name	Name of the direct child component

Returns: Shared pointer to component with correct type

Exceptions

std::out_of_range	if component name not found
std::runtime_error	if dynamic cast fails (type mismatch)

◆ getComponents()

template<DeviceType TDeviceType, TensorDataType TPrecision>

const std::vector< ComponentPtr > & Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::getComponents ( ) const

inline

Get all child components in insertion order.

Returns: Vector of child component pointers

◆ getDeviceId()

template<DeviceType TDeviceType, TensorDataType TPrecision>

DeviceId Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::getDeviceId ( ) const

inlineoverridevirtual

Get the compute device for this composite.

Returns the device from the shared execution context.

Returns: DeviceId for this composite and its children

Implements Mila::Dnn::Component< TDeviceType, TPrecision >.

Reimplemented in Mila::Dnn::Network< TDeviceType, TPrecision >.

◆ getGradients()

template<DeviceType TDeviceType, TensorDataType TPrecision>

std::vector< ITensor * > Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::getGradients ( ) const

inlineoverridevirtual

Get all parameter gradients from all children.

Gradient buffers are allocated only when built in training mode, so a component built for inference yields an empty result. This mirrors the leaf-component contract (e.g. Linear) and the symmetry with getParameters().

Returns: Vector of non-owning pointers to gradient tensors (empty in inference mode)

Exceptions

std::runtime_error if called before build()

Implements Mila::Dnn::Component< TDeviceType, TPrecision >.

◆ getParameters()

template<DeviceType TDeviceType, TensorDataType TPrecision>

std::vector< ITensor * > Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::getParameters ( ) const

inlineoverridevirtual

Get all parameters from all children.

Returns: Vector of non-owning pointers to parameter tensors

Exceptions

std::runtime_error if called before build()

Implements Mila::Dnn::Component< TDeviceType, TPrecision >.

◆ hasChildren()

template<DeviceType TDeviceType, TensorDataType TPrecision>

bool Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::hasChildren ( ) const

inlinenoexcept

Check if this composite has any children.

Returns: true if at least one child component exists

◆ hasComponent()

template<DeviceType TDeviceType, TensorDataType TPrecision>

bool Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::hasComponent ( const std::string & name ) const

inline

Check if a named child component exists.

Parameters

name	Name to query

Returns: true if a child with this name exists

◆ load_()

template<DeviceType TDeviceType, TensorDataType TPrecision>

void Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::load_	(	ModelArchive &	archive,
		SerializationMode	mode )

inlineoverrideprotectedvirtual

Restore children from their nested scopes, mirroring save_().

Walks the same children in the same order under the same scopes, so a child's blobs are read from exactly the paths they were written to. Restores into the live graph – the composite and its children must already be constructed and built.

Parameters

archive	Archive to read from, scoped to this composite.
mode	Serialization mode (passed to children).

Exceptions

std::runtime_error if the composite is not built.

Reimplemented from Mila::Dnn::Component< TDeviceType, TPrecision >.

Reimplemented in Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy >, Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, false, TWeightQuantization, TKvCachePolicy >, Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, true, TWeightQuantization, NoKvCompression >, Mila::Dnn::GptTransformer< TDeviceType, TPrecision >, and Mila::Dnn::Network< TDeviceType, TPrecision >.

◆ onExecutionContextSet()

template<DeviceType TDeviceType, TensorDataType TPrecision>

void Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::onExecutionContextSet ( )

inlineoverrideprotectedvirtual

Hook invoked after ExecutionContext is set.

Propagates the execution context to all child components that don't already have one. This enables the pattern where composites define their architecture graph in the constructor (context-independent) and context is bound later when available.

Called by Component::setExecutionContext() after the context is registered. Automatically invoked for both standalone mode (component creates own context) and shared mode (parent provides context).

Override this in derived classes if additional context-dependent initialization is required beyond context propagation to children.

Reimplemented from Mila::Dnn::Component< TDeviceType, TPrecision >.

◆ onTrainingModeChanging()

template<DeviceType TDeviceType, TensorDataType TPrecision>

void Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::onTrainingModeChanging ( TrainingMode training_mode )

inlineoverrideprotectedvirtual

Hook invoked when training mode is about to change.

Propagates the new mode to all child components. The hook runs with the Component's training mutex held; it MUST NOT call setTrainingMode().

Parameters

training_mode New training mode (Normal or Eval)

Reimplemented from Mila::Dnn::Component< TDeviceType, TPrecision >.

◆ optimize()

template<DeviceType TDeviceType, TensorDataType TPrecision>

virtual void Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::optimize ( )

inlineprotectedvirtual

Virtual hook for graph optimization after construction.

Called automatically after createGraph() completes. Derived classes can override to perform fusion, pruning, or other optimizations.

Default implementation does nothing. Override to perform architecture-specific graph optimizations.

◆ parameterCount()

template<DeviceType TDeviceType, TensorDataType TPrecision>

dim_t Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::parameterCount ( ) const

inlineoverridevirtual

Count parameters across all children.

Returns: Total number of trainable parameters

Exceptions

std::runtime_error if called before build()

Implements Mila::Dnn::Component< TDeviceType, TPrecision >.

Reimplemented in Mila::Dnn::GemmaTransformer< TDeviceType, TPrecision, TWeightQuantization, TKvCachePolicy >.

◆ removeComponent()

template<DeviceType TDeviceType, TensorDataType TPrecision>

bool Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::removeComponent ( const std::string & name )

inline

         @brief Get the named child components map.

         @return Map of names to child component pointers
        &zwj;/

const std::unordered_map<std::string, ComponentPtr>& getNamedComponents() const { return child_component_map_; }

    /**
       @brief Remove a child component by name.

       @param name Name of the component to remove
       @return true if removed, false if not found

       @throws std::runtime_error if called after build()

◆ requireSerializableParameters()

template<DeviceType TDeviceType, TensorDataType TPrecision>

void Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::requireSerializableParameters ( ) const

inlineoverrideprotectedvirtual

No-op override: a composite names no parameters of its own.

parameterCount() on a composite sums its children, so the base implementation would demand names this component never owns. Each child is checked individually by the recursion in save_().

Reimplemented from Mila::Dnn::Component< TDeviceType, TPrecision >.

◆ save_()

template<DeviceType TDeviceType, TensorDataType TPrecision>

void Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::save_	(	ModelArchive &	archive,
		SerializationMode	mode ) const

inlineoverrideprotectedvirtual

Implements Mila::Dnn::Component< TDeviceType, TPrecision >.

Reimplemented in Mila::Dnn::GemmaTransformer< TDeviceType, TPrecision, TWeightQuantization, TKvCachePolicy >, Mila::Dnn::GptTransformer< TDeviceType, TPrecision >, Mila::Dnn::LlamaTransformer< TDeviceType, TPrecision, TWeightQuantization, TKvCachePolicy >, and Mila::Dnn::Network< TDeviceType, TPrecision >.

◆ saveFlatTensors()

template<DeviceType TDeviceType, TensorDataType TPrecision>

void Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::saveFlatTensors	(	Serialization::SafeTensorsWriter &	writer,
		const std::string &	prefix,
		Serialization::TensorSavePass	pass ) const

inlineoverridevirtual

Recurse into children, extending the flat dotted prefix.

A composite contributes no tensors of its own. It exists here to turn the component tree into the flat vocabulary the pretrained format uses – the same dotted paths loadParameters() splits with parseParameterPath() and resolves with findComponent(), so what this writes is exactly what that reads.

Children carry fully qualified names ("gemma.tf_layer_0.qkv_proj"), so the relative segment is recovered by stripping this composite's own name. Calling the root with an empty prefix therefore drops the model name and yields "tf_layer_0.qkv_proj", matching the converter's naming rather than a Mila-internal path.

child_components_ rather than child_component_map_: the map is unordered, and the writer requires bodies in declaration order, so an unordered walk would produce a file whose data region disagrees with its own index.

Public, unlike save_() which is public on Component but protected here – an asymmetry already filed as a defect. A flat save is driven from outside the tree, so repeating it would force every caller through a forwarder.

Reimplemented from Mila::Dnn::Component< TDeviceType, TPrecision >.

Reimplemented in Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy >, Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, false, TWeightQuantization, TKvCachePolicy >, and Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, true, TWeightQuantization, NoKvCompression >.

◆ synchronize()

template<DeviceType TDeviceType, TensorDataType TPrecision>

void Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::synchronize ( )

inlineoverridevirtual

Synchronize all child components.

Waits for outstanding device operations on all children.

Implements Mila::Dnn::Component< TDeviceType, TPrecision >.

Reimplemented in Mila::Dnn::Network< TDeviceType, TPrecision >.

◆ toString()

template<DeviceType TDeviceType, TensorDataType TPrecision>

std::string Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::toString ( ) const

inlineoverridevirtual

Generate a human-readable description.

Returns: String representation showing children

Implements Mila::Dnn::Component< TDeviceType, TPrecision >.

Reimplemented in Mila::Dnn::GatedMLP< TDeviceType, TPrecision, TGate >, Mila::Dnn::GemmaTransformer< TDeviceType, TPrecision, TWeightQuantization, TKvCachePolicy >, Mila::Dnn::GptBlock< TDeviceType, TPrecision >, Mila::Dnn::GptTransformer< TDeviceType, TPrecision >, Mila::Dnn::LlamaTransformer< TDeviceType, TPrecision, TWeightQuantization, TKvCachePolicy >, Mila::Dnn::MLP< TDeviceType, TPrecision >, and Mila::Dnn::Network< TDeviceType, TPrecision >.

◆ tryFindComponent()

template<DeviceType TDeviceType, TensorDataType TPrecision>

ComponentPtr Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >::tryFindComponent ( const std::string & path ) const

inline

Try to resolve a dot-separated component path within this composite.

Non-throwing version: returns nullptr if any segment is not found or if a path segment attempts to traverse into a non-composite leaf.

Example: auto ptr = composite->tryFindComponent("encoder.mlp.fc1");

Parameters

path	Dot-separated path (e.g. "layer_0.mlp.fc_1")

Returns: ComponentPtr or nullptr when not found / invalid traversal

The documentation for this class was generated from the following file:

Mila/Src/Dnn/Core/CompositeComponent.ixx

Public Types

Public Member Functions

Protected Member Functions

Friends

Additional Inherited Members

Detailed Description

Constructor & Destructor Documentation

◆ CompositeComponent()

Member Function Documentation

◆ addComponent()

◆ childCount()

◆ clearComponents()

◆ findComponent()

◆ getComponent()

◆ getComponentAs()

◆ getComponents()

◆ getDeviceId()

◆ getGradients()

◆ getParameters()

◆ hasChildren()

◆ hasComponent()

◆ load_()

◆ onExecutionContextSet()

◆ onTrainingModeChanging()

◆ optimize()

◆ parameterCount()

◆ removeComponent()

◆ requireSerializableParameters()

◆ save_()

◆ saveFlatTensors()

◆ synchronize()

◆ toString()

◆ tryFindComponent()