Mila::Dnn::Compute::Operation< TDeviceType, TComputePrecision > Class Template Reference

Public Types
using	DataTypeTraits = TensorDataTypeTraits<TComputePrecision>

Public Member Functions
virtual void	build (const BuildContext &build_context)
	Prepare the operation for a concrete input shape.
virtual void	clearGradients () noexcept
	Clear any cached gradient pointers held by the operation.
virtual TensorDataType	getDataType () const
	Tensor data type for this operation.
virtual DeviceType	getDeviceType () const
	Device type for this operation.
virtual std::string	getName () const =0
	Human-readable operation name.
virtual OperationType	getOperationType () const =0
	Operation type identifier.
virtual std::size_t	getRequiredStateMemorySize (const BuildContext &build_context) const
	Bytes of state build( context ) would allocate, without allocating it.
virtual std::size_t	getStateMemorySize () const
	Returns the number of bytes of state memory allocated by this operation.
virtual bool	isBuilt () const
	Whether build() completed successfully for a concrete input shape.
virtual bool	isEvalMode () const
	Query whether operation is configured for training.
virtual void	setGradients (ITensor *weight_grad, ITensor *bias_grad)
	Bind module-owned gradient tensors to the operation.
virtual void	setParameters (ITensor *weight, ITensor *bias)
	Bind module-owned parameter tensors to the operation.
virtual void	setTrainingMode (TrainingMode training_mode)
	Configure operation training-mode behavior.

Static Public Attributes
static constexpr TensorDataType	data_type = TComputePrecision
static constexpr DeviceType	device_type = TDeviceType

Protected Attributes
bool	is_built_ { false }
TrainingMode	training_mode_ { TrainingMode::Normal }

Member Function Documentation

build()

template<DeviceType TDeviceType, TensorDataType TComputePrecision>

virtual void Mila::Dnn::Compute::Operation< TDeviceType, TComputePrecision >::build ( const BuildContext & build_context )

inlinevirtual

Prepare the operation for a concrete input shape.

Default implementation is a no-op. Operations requiring shape-dependent setup should override this method.

clearGradients()

template<DeviceType TDeviceType, TensorDataType TComputePrecision>

virtual void Mila::Dnn::Compute::Operation< TDeviceType, TComputePrecision >::clearGradients ( )

inlinevirtualnoexcept

Clear any cached gradient pointers held by the operation.

Explicit unbind called by modules before freeing/resetting module-owned gradient buffers. Implementations MUST null-out any cached raw pointers and MUST NOT throw. Marked noexcept so it is safe to call from destructors or during state transitions.

getRequiredStateMemorySize()

template<DeviceType TDeviceType, TensorDataType TComputePrecision>

virtual std::size_t Mila::Dnn::Compute::Operation< TDeviceType, TComputePrecision >::getRequiredStateMemorySize ( const BuildContext & build_context ) const

inlinevirtual

Bytes of state build( context ) would allocate, without allocating it.

The predictive counterpart to getStateMemorySize(). Zero by default because an operation allocates no state unless it says otherwise – the same premise getStateMemorySize() already defaults on, so the two agree for every stateless operation without either being written.

Override wherever getStateMemorySize() is overridden. The KV cache capacity rule belongs here rather than in the component: the cache is allocated on this side of the boundary.

getStateMemorySize()

template<DeviceType TDeviceType, TensorDataType TComputePrecision>

virtual std::size_t Mila::Dnn::Compute::Operation< TDeviceType, TComputePrecision >::getStateMemorySize ( ) const

inlinevirtual

Returns the number of bytes of state memory allocated by this operation.

State memory includes build-time buffers such as caches and scratch allocations. Parameters and gradients are owned at the component level and are not included.

Override in derived operations that allocate device or host state during build().

setGradients()

template<DeviceType TDeviceType, TensorDataType TComputePrecision>

virtual void Mila::Dnn::Compute::Operation< TDeviceType, TComputePrecision >::setGradients ( ITensor * weight_grad, ITensor * bias_grad )

inlinevirtual

Bind module-owned gradient tensors to the operation.

New canonical API for binding gradient buffers. Mirrors semantics of setParameters() but for gradients used during backward().

The operation MUST NOT take ownership of the provided pointers. Implementations may cache rawData() pointers for hot-path writes.

Default: no-op for stateless operations.

setParameters()

template<DeviceType TDeviceType, TensorDataType TComputePrecision>

virtual void Mila::Dnn::Compute::Operation< TDeviceType, TComputePrecision >::setParameters ( ITensor * weight, ITensor * bias )

inlinevirtual

Bind module-owned parameter tensors to the operation.

The module retains ownership of the provided ITensor objects. Implementations may cache rawData() pointers for hot-path access but MUST NOT free the provided pointers.

Default: no-op for stateless operations.

setTrainingMode()

template<DeviceType TDeviceType, TensorDataType TComputePrecision>

virtual void Mila::Dnn::Compute::Operation< TDeviceType, TComputePrecision >::setTrainingMode ( TrainingMode training_mode )

inlinevirtual

Configure operation training-mode behavior.

Implementations may use this to enable/disable training-specific work.

---

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

• Mila/Src/Dnn/Compute/Operations/OperationBase.ixx