Softmax activation module (device-templated). More...

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

Public Types
using	ComponentBase = Component<TDeviceType, TPrecision>
using	MR = typename DeviceTypeTraits<TDeviceType>::memory_resource
using	TensorType = Tensor<TPrecision, MR>

Public Member Functions
	Softmax (const std::string &name, const SoftmaxConfig &config, std::optional< DeviceId > device_id=std::nullopt)
	Construct Softmax with optional ExecutionContext ownership.
void	backward (const ITensor &input, const ITensor &output_grad, ITensor &input_grad)
	Backward pass - delegates to backend operation.
void	forward (const ITensor &input, ITensor &output)
	Forward pass - delegates to backend operation.
int64_t	getAxis () const noexcept
	Get the softmax axis.
DeviceId	getDeviceId () const override
	Get the device identifier for this module.
std::vector< ITensor * >	getGradients () const override
	Get parameter gradient tensors.
MemoryStats	getMemoryStats () const override
	Return the current memory allocation breakdown for this component.
std::vector< ITensor * >	getParameters () const override
	Get trainable parameter tensors.
const ComponentType	getType () const override
	Get the component type identifier.
dim_t	parameterCount () const override
	Number of trainable parameters.
void	save_ (ModelArchive &, SerializationMode) const override
	Persist module state to archive.
void	synchronize () override
	Wait for all asynchronous work submitted by this module to complete.
std::string	toString () const override
	Generate human-readable description of the module.
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).
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 bool	isBuilt () const final
	Returns true if build() has completed successfully.
virtual void	load_ (ModelArchive &archive, SerializationMode mode)
	Restore this component's parameters from its archive scope.
virtual void	loadParameter (const std::string &, const Serialization::ITensorBlob &)
	Load a parameter from serialized tensor data.
virtual void	requireSerializableParameters () const
	Verify this component can serialize whatever parameters it owns.
virtual void	saveFlatTensors (Serialization::SafeTensorsWriter &writer, const std::string &prefix, Serialization::TensorSavePass pass) const
	Drive this component's tensors through one pass of a flat safetensors save.
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
void	onBuilding (const BuildContext &build_config) override
	Hook invoked during build() to initialize component with input shape.
void	onExecutionContextSet () override
	Hook invoked after ExecutionContext is set.
void	onTrainingModeChanging (TrainingMode training_mode) override
	Hook invoked when training mode changes.
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.
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.

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>
requires PrecisionSupportedOnDevice<TPrecision, TDeviceType>
class Mila::Dnn::Softmax< TDeviceType, TPrecision >

Softmax activation module (device-templated).

Delegates computation to a device-specific UnaryOperation implementation registered in the OperationRegistry.

Softmax is a stateless activation function with no trainable parameters. The operation computes: softmax(x) = exp(x - max(x)) / sum(exp(x - max(x))) across a specified axis.

Construction Modes:

Standalone mode: Construct with DeviceId to create and own an ExecutionContext. The component manages the context lifetime and uses it for operation execution.
Shared mode: Construct without DeviceId; parent (Network/CompositeComponent) provides ExecutionContext via setExecutionContext() after construction.

Ownership:

Standalone mode: Component owns its ExecutionContext (stored in owned_exec_context_).
Shared mode: Component borrows ExecutionContext from parent; lifecycle managed externally.

Template Parameters

TDeviceType	Device type (DeviceType::Cpu or DeviceType::Cuda)
TPrecision	Abstract tensor precision (TensorDataType)

Constructor & Destructor Documentation

◆ Softmax()

template<DeviceType TDeviceType, TensorDataType TPrecision>

Mila::Dnn::Softmax< TDeviceType, TPrecision >::Softmax	(	const std::string &	name,
		const SoftmaxConfig &	config,
		std::optional< DeviceId >	device_id = std::nullopt )

inlineexplicit

Construct Softmax with optional ExecutionContext ownership.

Supports two construction modes:

Standalone mode (device_id provided):

Creates and owns an ExecutionContext for the specified device.
Registers the owned context with the base Component class via setExecutionContext().
Backend operation is created immediately in onExecutionContextSet() hook.
Use case: Unit tests, standalone component usage.

Shared mode (device_id not provided):

Does not create ExecutionContext; expects parent to provide one.
Parent (Network/CompositeComponent) calls setExecutionContext() after construction.
Backend operation created when parent sets context.
Use case: Components added to Network via addComponent<Softmax>(...).

Parameters

name	Component name identifier.
config	Softmax configuration (axis and name).
device_id	Optional device identifier. If provided, creates owned ExecutionContext for standalone mode. If nullopt, expects shared context from parent.

Exceptions

std::invalid_argument	if config is invalid (via config.validate()).
std::invalid_argument	if device_id.type does not match TDeviceType.
std::runtime_error	if ExecutionContext creation fails (standalone mode).
std::runtime_error	if backend operation creation fails in onExecutionContextSet().

Note: In standalone mode, setExecutionContext() is called to register the owned context with the base class, enabling getExecutionContext() and triggering the onExecutionContextSet() hook for operation creation.

// Standalone mode (owns context)
SoftmaxConfig config;
config.withAxis(-1);
Softmax<DeviceType::Cpu, TensorDataType::FP32> softmax(config, Device::Cpu());

// Shared mode (borrows parent's context)
Network<DeviceType::Cpu, TensorDataType::FP32> net(Device::Cpu(), "my_net");
net.addComponent<Softmax>("softmax", SoftmaxConfig().withAxis(-1));

Member Function Documentation

◆ backward()

template<DeviceType TDeviceType, TensorDataType TPrecision>

void Mila::Dnn::Softmax< TDeviceType, TPrecision >::backward	(	const ITensor &	input,
		const ITensor &	output_grad,
		ITensor &	input_grad )

inline

Backward pass - delegates to backend operation.

Computes gradient: dX = Y * (dY - dot(Y, dY)) where Y is the softmax output.

◆ forward()

template<DeviceType TDeviceType, TensorDataType TPrecision>

void Mila::Dnn::Softmax< TDeviceType, TPrecision >::forward	(	const ITensor &	input,
		ITensor &	output )

inline

Forward pass - delegates to backend operation.

Computes softmax activation across the configured axis.

◆ getAxis()

template<DeviceType TDeviceType, TensorDataType TPrecision>

int64_t Mila::Dnn::Softmax< TDeviceType, TPrecision >::getAxis ( ) const

inlinenoexcept

Get the softmax axis.

Returns: The axis along which softmax is computed.

◆ getDeviceId()

template<DeviceType TDeviceType, TensorDataType TPrecision>

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

inlineoverridevirtual

Get the device identifier for this module.

Returns the DeviceId from the ExecutionContext. In standalone mode, this is the device specified at construction. In shared mode, this is the parent's device.

Returns: DeviceId indicating device type and index.

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

◆ getGradients()

template<DeviceType TDeviceType, TensorDataType TPrecision>

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

inlineoverridevirtual

Get parameter gradient tensors.

Softmax has no trainable parameters, therefore no gradients.

Returns: Empty vector.

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

◆ getMemoryStats()

template<DeviceType TDeviceType, TensorDataType TPrecision>

MemoryStats Mila::Dnn::Softmax< TDeviceType, TPrecision >::getMemoryStats ( ) const

inlineoverridevirtual

Return the current memory allocation breakdown for this component.

Reflects allocations at the moment of the call. The returned stats naturally track the component lifecycle:

After construction – nothing; construction allocates none After build( Inference ) – parameters + T=1 state buffers After build( Training ) – parameters + T=full state buffers After setTrainingMode( Train ) – parameters + state + gradients

For CompositeComponent and Network, the returned stats are the recursive aggregate of all child components.

May be called at any time – no lifecycle preconditions.

Returns: MemoryStats reflecting current allocations.

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

◆ getParameters()

template<DeviceType TDeviceType, TensorDataType TPrecision>

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

inlineoverridevirtual

Get trainable parameter tensors.

Softmax has no trainable parameters.

Returns: Empty vector.

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

◆ getType()

template<DeviceType TDeviceType, TensorDataType TPrecision>

const ComponentType Mila::Dnn::Softmax< TDeviceType, TPrecision >::getType ( ) const

inlineoverridevirtual

Get the component type identifier.

Used for serialization and runtime type identification.

Returns: Component type enum value.

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

◆ onBuilding()

template<DeviceType TDeviceType, TensorDataType TPrecision>

void Mila::Dnn::Softmax< TDeviceType, TPrecision >::onBuilding ( const BuildContext & build_config )

inlineoverrideprotectedvirtual

Hook invoked during build() to initialize component with input shape.

Softmax is stateless and has no parameters to allocate. This method validates the input shape and delegates to the backend operation's build method to cache dimension computations.

Parameters

build_config Build-time context carrying the expected input shape.

Exceptions

std::invalid_argument	if the input shape is invalid or axis out of bounds.
std::runtime_error	if backend build fails.

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

◆ onExecutionContextSet()

template<DeviceType TDeviceType, TensorDataType TPrecision>

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

inlineoverrideprotectedvirtual

Hook invoked after ExecutionContext is set.

Called by Component::setExecutionContext() after the context is registered. Creates the backend UnaryOperation using the OperationRegistry.

This hook is triggered in two scenarios:

Standalone mode: Immediately in constructor after owned context creation
Shared mode: When parent calls setExecutionContext() after construction

Exceptions

std::runtime_error if operation creation fails.

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

◆ onTrainingModeChanging()

template<DeviceType TDeviceType, TensorDataType TPrecision>

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

inlineoverrideprotectedvirtual

Hook invoked when training mode changes.

Propagates training mode to the backend operation. Called by Component::setTrainingMode() with the training mutex held.

Parameters

training_mode New training mode state.

Note: Do not call setTrainingMode() from this hook (reentrancy prohibited).

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

◆ parameterCount()

template<DeviceType TDeviceType, TensorDataType TPrecision>

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

inlineoverridevirtual

Number of trainable parameters.

Softmax is stateless and exposes no trainable parameters.

Returns: 0

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

◆ save_()

template<DeviceType TDeviceType, TensorDataType TPrecision>

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

inlineoverridevirtual

Persist module state to archive.

Softmax is stateless (no trainable tensors) but persists:

Module type and version metadata
Configuration (axis)

Takes the archive to write to, and a serialization mode that is currently unused for stateless components.

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

◆ synchronize()

template<DeviceType TDeviceType, TensorDataType TPrecision>

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

inlineoverridevirtual

Wait for all asynchronous work submitted by this module to complete.

Synchronizes the underlying ExecutionContext. On CPU implementations this may be a no-op. Use to ensure results are visible on the host or to measure synchronous timings.

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

◆ toString()

template<DeviceType TDeviceType, TensorDataType TPrecision>

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

inlineoverridevirtual

Generate human-readable description of the module.

Produces a multi-line string showing:

Module name
Device type
Axis configuration

Returns: Formatted string representation.

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

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

Mila/Src/Dnn/Components/Normalization/Softmax.ixx

Public Types

Public Member Functions

Protected Member Functions

Additional Inherited Members

Detailed Description

Constructor & Destructor Documentation

◆ Softmax()

Member Function Documentation

◆ backward()

◆ forward()

◆ getAxis()

◆ getDeviceId()

◆ getGradients()

◆ getMemoryStats()

◆ getParameters()

◆ getType()

◆ onBuilding()

◆ onExecutionContextSet()

◆ onTrainingModeChanging()

◆ parameterCount()

◆ save_()

◆ synchronize()

◆ toString()