Fused SoftmaxCrossEntropy loss module (device-templated). More...

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

Public Types
using	MR = typename DeviceTypeTraits<TDeviceType>::memory_resource
using	TargetTensorType = Tensor<TensorDataType::INT32, MR>
using	TensorType = Tensor<TPrecision, MR>

Public Member Functions
	SoftmaxCrossEntropy (IExecutionContext *exec_context, const CrossEntropyConfig &config)
	Construct with an existing execution context.
void	backward (const ITensor &logits, const ITensor &targets, const ITensor &output_grad, ITensor &logits_grad)
	Backward pass - delegates to backend operation.
void	forward (const ITensor &logits, const ITensor &targets, ITensor &output)
	Forward pass - delegates to backend operation.
const CrossEntropyConfig &	getConfig () const noexcept
DeviceId	getDeviceId () const override
	Get the compute device id associated with this component.
std::vector< ITensor * >	getGradients () const override
	Return non-owning pointers to parameter gradient tensors.
std::vector< ITensor * >	getParameters () const override
	Return non-owning pointers to parameter tensors.
int64_t	getVocabSize () const
void	onBuilding (const BuildContext &context) override
	Build the module using an input shape.
dim_t	parameterCount () const override
	Return number of trainable parameters.
void	save_ (ModelArchive &, SerializationMode) const override
void	synchronize () override
	Wait for outstanding device work submitted by this component.
std::string	toString () const override
	Produce a short, human-readable description of the component.
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).
IExecutionContext *	getExecutionContext () const
	Get the shared execution context.
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< ObservableStage >	getObservableStages () const
	Stages this component publishes, and the passes each is published on.
Compute::ComputePassMask	getObservedPasses () const noexcept
virtual std::vector< const ITensor * >	getOutputs () const
	Return non-owning pointers to the tensors this component produces.
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.
RuntimeMode	getRuntimeMode () const noexcept
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.
bool	isInferenceMode () const noexcept
bool	isTrainingMode () const noexcept
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	setObservedPasses (Compute::ComputePassMask passes) noexcept
	Select which passes this component publishes on.
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	onExecutionContextSet () override
	Hook invoked after the ExecutionContext is set.
void	onTrainingModeChanging (TrainingMode training_mode) override
	Hook invoked when training mode is about to change.
Protected Member Functions inherited from Mila::Dnn::Component< TDeviceType, TPrecision >
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.
void	publish (Compute::ComputePass pass, std::string_view stage, const ITensor &value) const
	Publish one activation to the observer, if this pass is being observed.
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 = typename DeviceTypeTraits<TDeviceType>::host_staging_memory_resource
	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::SoftmaxCrossEntropy< TDeviceType, TPrecision >

Fused SoftmaxCrossEntropy loss module (device-templated).

Delegates computation to a device-specific operation resolved at compile time via OperationTraits<CrossEntropyOp, TDeviceType, TPrecision>. Targets are discrete class indices (INT32) and are not a precision axis.

Constructor & Destructor Documentation

◆ SoftmaxCrossEntropy()

template<DeviceType TDeviceType, TensorDataType TPrecision>

Mila::Dnn::SoftmaxCrossEntropy< TDeviceType, TPrecision >::SoftmaxCrossEntropy	(	IExecutionContext *	exec_context,
		const CrossEntropyConfig &	config )

inlineexplicit

Construct with an existing execution context.

Parameters

exec_context	Shared execution context for device resources.
config	CrossEntropy configuration (vocab_size required).

Member Function Documentation

◆ backward()

template<DeviceType TDeviceType, TensorDataType TPrecision>

void Mila::Dnn::SoftmaxCrossEntropy< TDeviceType, TPrecision >::backward	(	const ITensor &	logits,
		const ITensor &	targets,
		const ITensor &	output_grad,
		ITensor &	logits_grad )

inline

Backward pass - delegates to backend operation.

Computes fused gradient: dL/dlogits = softmax(logits) - one_hot(targets)

◆ forward()

template<DeviceType TDeviceType, TensorDataType TPrecision>

void Mila::Dnn::SoftmaxCrossEntropy< TDeviceType, TPrecision >::forward	(	const ITensor &	logits,
		const ITensor &	targets,
		ITensor &	output )

inline

Forward pass - delegates to backend operation.

Computes fused softmax + cross-entropy loss.

◆ getDeviceId()

template<DeviceType TDeviceType, TensorDataType TPrecision>

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

inlineoverridevirtual

Get the compute device id associated with this component.

Must return the device on which parameters and operations execute.

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

◆ getGradients()

template<DeviceType TDeviceType, TensorDataType TPrecision>

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

inlineoverridevirtual

Return non-owning pointers to parameter gradient tensors.

Gradient buffers are allocated only when the component is built in training mode, so a component built for inference returns an empty vector. Stateless components return empty in either mode. This is the accessor counterpart to getParameters() and does not throw on mode.

Returns: Vector of gradient pointers; empty when built for inference.

Exceptions

std::runtime_error if called before the component has been built.

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

◆ getParameters()

template<DeviceType TDeviceType, TensorDataType TPrecision>

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

inlineoverridevirtual

Return non-owning pointers to parameter tensors.

The returned tensor pointers remain valid for the lifetime of the component. Order should be canonical (weights before biases).

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

◆ onBuilding()

template<DeviceType TDeviceType, TensorDataType TPrecision>

void Mila::Dnn::SoftmaxCrossEntropy< TDeviceType, TPrecision >::onBuilding ( const BuildContext & context )

inlineoverridevirtual

Build the module using an input shape.

Validates input shape and triggers backend-specific setup. The fused operation has no trainable parameters.

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

◆ onExecutionContextSet()

template<DeviceType TDeviceType, TensorDataType TPrecision>

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

inlineoverrideprotectedvirtual

Hook invoked after the ExecutionContext is set.

Creates the backend operation via compile-time OperationTraits dispatch.

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

◆ onTrainingModeChanging()

template<DeviceType TDeviceType, TensorDataType TPrecision>

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

inlineoverrideprotectedvirtual

Hook invoked when training mode is about to change.

Propagate training mode to the backend fused operation. Called with the training mutex held; do not call setTrainingMode() on the component here.

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

◆ parameterCount()

template<DeviceType TDeviceType, TensorDataType TPrecision>

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

inlineoverridevirtual

Return number of trainable parameters.

For leaf components this is the element count of owned parameter tensors. CompositeComponent and Network implementations should return the recursive aggregate across all children.

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

◆ save_()

template<DeviceType TDeviceType, TensorDataType TPrecision>

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

inlineoverridevirtual

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

◆ synchronize()

template<DeviceType TDeviceType, TensorDataType TPrecision>

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

inlineoverridevirtual

Wait for outstanding device work submitted by this component.

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

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

◆ toString()

template<DeviceType TDeviceType, TensorDataType TPrecision>

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

inlineoverridevirtual

Produce a short, human-readable description of the component.

Implementations should keep output concise and avoid throwing.

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

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

Mila/Src/Dnn/Components/Losses/SoftmaxCrossEntropy.ixx

Public Types

Public Member Functions

Protected Member Functions

Additional Inherited Members

Detailed Description

Constructor & Destructor Documentation

◆ SoftmaxCrossEntropy()

Member Function Documentation

◆ backward()

◆ forward()

◆ getDeviceId()

◆ getGradients()

◆ getParameters()

◆ onBuilding()

◆ onExecutionContextSet()

◆ onTrainingModeChanging()

◆ parameterCount()

◆ save_()

◆ synchronize()

◆ toString()