Mila 0.13.48
Deep Neural Network Library
Loading...
Searching...
No Matches
Mila::Dnn::Compute::UnaryOperation< TDeviceType, TInput, TPrecision > Class Template Referenceabstractexport
module Compute.UnaryOperation
Inheritance diagram for Mila::Dnn::Compute::UnaryOperation< TDeviceType, TInput, TPrecision >:
Inheritance graph
[legend]
Collaboration diagram for Mila::Dnn::Compute::UnaryOperation< TDeviceType, TInput, TPrecision >:
Collaboration graph
[legend]

Public Types

using MR = CpuMemoryResource
using TensorInputType = Tensor<TInput, MR>
using TensorOutputType = Tensor<TPrecision, MR>
Public Types inherited from Mila::Dnn::Compute::Operation< TDeviceType, TInput >
using DataTypeTraits

Public Member Functions

virtual ~UnaryOperation ()=default
virtual void backward (const ITensor &input, const ITensor &output_grad, ITensor &input_grad) const =0
 Backward pass: compute gradient wrt input given output gradient.
virtual void forward (const ITensor &input, ITensor &output) const =0
 Forward pass: compute output = f(input).
Public Member Functions inherited from Mila::Dnn::Compute::Operation< TDeviceType, TInput >
virtual ~Operation ()=default
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 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 Protected Member Functions

static const TensorInputType & asInputTensor (const ITensor &t)
static TensorOutputType & asOutputTensor (ITensor &t)

Additional Inherited Members

Static Public Attributes inherited from Mila::Dnn::Compute::Operation< TDeviceType, TInput >
static constexpr TensorDataType data_type
static constexpr DeviceType device_type
Protected Attributes inherited from Mila::Dnn::Compute::Operation< TDeviceType, TInput >
bool is_built_
TrainingMode training_mode_

Detailed Description

template<DeviceType TDeviceType, TensorDataType TInput, TensorDataType TPrecision = TInput>
requires PrecisionSupportedOnDevice<TPrecision, TDeviceType>
class Mila::Dnn::Compute::UnaryOperation< TDeviceType, TInput, TPrecision >
Template Parameters
TDeviceTypeDevice target for the operation (DeviceType::Cpu, DeviceType::Cuda, ...)
TPrecisionCanonical element precision produced/consumed by the op (e.g. FP32)
TInputOptional element type for the runtime input tensor (defaults to TPrecision, e.g. INT32 for token indices while TPrecision == FP32)
Note
Constrains are applied to ensure the operation precision is supported on the device.

Member Typedef Documentation

◆ MR

template<DeviceType TDeviceType, TensorDataType TInput, TensorDataType TPrecision = TInput>
using Mila::Dnn::Compute::UnaryOperation< TDeviceType, TInput, TPrecision >::MR = CpuMemoryResource

◆ TensorInputType

template<DeviceType TDeviceType, TensorDataType TInput, TensorDataType TPrecision = TInput>
using Mila::Dnn::Compute::UnaryOperation< TDeviceType, TInput, TPrecision >::TensorInputType = Tensor<TInput, MR>

◆ TensorOutputType

template<DeviceType TDeviceType, TensorDataType TInput, TensorDataType TPrecision = TInput>
using Mila::Dnn::Compute::UnaryOperation< TDeviceType, TInput, TPrecision >::TensorOutputType = Tensor<TPrecision, MR>

Constructor & Destructor Documentation

◆ ~UnaryOperation()

template<DeviceType TDeviceType, TensorDataType TInput, TensorDataType TPrecision = TInput>
virtual Mila::Dnn::Compute::UnaryOperation< TDeviceType, TInput, TPrecision >::~UnaryOperation ( )
virtualdefault

Member Function Documentation

◆ asInputTensor()

template<DeviceType TDeviceType, TensorDataType TInput, TensorDataType TPrecision = TInput>
const TensorInputType & Mila::Dnn::Compute::UnaryOperation< TDeviceType, TInput, TPrecision >::asInputTensor ( const ITensor & t)
inlinestaticprotected

◆ asOutputTensor()

template<DeviceType TDeviceType, TensorDataType TInput, TensorDataType TPrecision = TInput>
TensorOutputType & Mila::Dnn::Compute::UnaryOperation< TDeviceType, TInput, TPrecision >::asOutputTensor ( ITensor & t)
inlinestaticprotected

◆ backward()

template<DeviceType TDeviceType, TensorDataType TInput, TensorDataType TPrecision = TInput>
virtual void Mila::Dnn::Compute::UnaryOperation< TDeviceType, TInput, TPrecision >::backward ( const ITensor & input,
const ITensor & output_grad,
ITensor & input_grad ) const
pure virtual

Backward pass: compute gradient wrt input given output gradient.

Signature ordered as (input, output_grad, input_grad) to match module and operation implementations across the codebase.

Implemented in Mila::Dnn::Compute::CpuAttentionOp, Mila::Dnn::Compute::CpuEncoderOp, Mila::Dnn::Compute::CpuGeluOp, Mila::Dnn::Compute::CpuLayerNormOp, Mila::Dnn::Compute::CpuLinearOp, Mila::Dnn::Compute::CpuSoftmaxOp, Mila::Dnn::Compute::Cuda::Gelu::CudaGeluOp< TPrecision >, Mila::Dnn::Compute::Cuda::Gelu::CudaGeluOp< TensorDataType::BF16 >, Mila::Dnn::Compute::Cuda::Gelu::CudaGeluOp< TensorDataType::FP32 >, Mila::Dnn::Compute::Cuda::LayerNorm::CudaLayerNormOp< TPrecision >, Mila::Dnn::Compute::Cuda::Lpe::CudaLpeOp< TInput, TPrecision >, Mila::Dnn::Compute::Cuda::Lpe::CudaLpeOp< TensorDataType::INT32, TensorDataType::BF16 >, Mila::Dnn::Compute::Cuda::Lpe::CudaLpeOp< TensorDataType::INT32, TensorDataType::FP32 >, Mila::Dnn::Compute::Cuda::MultiHeadAttention::CudaMultiHeadAttentionOp< TPrecision >, Mila::Dnn::Compute::Cuda::MultiHeadAttention::CudaMultiHeadAttentionOp< TensorDataType::BF16 >, Mila::Dnn::Compute::Cuda::MultiHeadAttention::CudaMultiHeadAttentionOp< TensorDataType::FP32 >, Mila::Dnn::Compute::Cuda::RmsNorm::CudaRmsNormOp< TPrecision >, Mila::Dnn::Compute::Cuda::RmsNorm::CudaRmsNormOp< TensorDataType::BF16 >, Mila::Dnn::Compute::Cuda::RmsNorm::CudaRmsNormOp< TensorDataType::FP32 >, Mila::Dnn::Compute::Cuda::Softmax::CudaSoftmaxOp< TPrecision >, Mila::Dnn::Compute::Cuda::Softmax::CudaSoftmaxOp< TensorDataType::BF16 >, Mila::Dnn::Compute::Cuda::Softmax::CudaSoftmaxOp< TensorDataType::FP32 >, Mila::Dnn::Compute::Cuda::Swiglu::CudaSwigluOp< TPrecision >, Mila::Dnn::Compute::Cuda::Swiglu::CudaSwigluOp< TensorDataType::BF16 >, Mila::Dnn::Compute::Cuda::Swiglu::CudaSwigluOp< TensorDataType::FP32 >, Mila::Dnn::Compute::Cuda::TokenEmbedding::CudaTokenEmbeddingOp< TInput, TPrecision >, Mila::Dnn::Compute::Cuda::TokenEmbedding::CudaTokenEmbeddingOp< TensorDataType::INT32, TensorDataType::BF16 >, and Mila::Dnn::Compute::Cuda::TokenEmbedding::CudaTokenEmbeddingOp< TensorDataType::INT32, TensorDataType::FP32 >.

◆ forward()

template<DeviceType TDeviceType, TensorDataType TInput, TensorDataType TPrecision = TInput>
virtual void Mila::Dnn::Compute::UnaryOperation< TDeviceType, TInput, TPrecision >::forward ( const ITensor & input,
ITensor & output ) const
pure virtual

Forward pass: compute output = f(input).

Implementations should accept polymorphic ITensor references and may use the typed aliases / helpers to obtain typed tensor references.

Implemented in Mila::Dnn::Compute::CpuAttentionOp, Mila::Dnn::Compute::CpuEncoderOp, Mila::Dnn::Compute::CpuGeluOp, Mila::Dnn::Compute::CpuLayerNormOp, Mila::Dnn::Compute::CpuLinearOp, Mila::Dnn::Compute::CpuSoftmaxOp, Mila::Dnn::Compute::Cuda::Gelu::CudaGeluOp< TPrecision >, Mila::Dnn::Compute::Cuda::Gelu::CudaGeluOp< TensorDataType::BF16 >, Mila::Dnn::Compute::Cuda::Gelu::CudaGeluOp< TensorDataType::FP32 >, Mila::Dnn::Compute::Cuda::LayerNorm::CudaLayerNormOp< TPrecision >, Mila::Dnn::Compute::Cuda::Lpe::CudaLpeOp< TInput, TPrecision >, Mila::Dnn::Compute::Cuda::Lpe::CudaLpeOp< TensorDataType::INT32, TensorDataType::BF16 >, Mila::Dnn::Compute::Cuda::Lpe::CudaLpeOp< TensorDataType::INT32, TensorDataType::FP32 >, Mila::Dnn::Compute::Cuda::MultiHeadAttention::CudaMultiHeadAttentionOp< TPrecision >, Mila::Dnn::Compute::Cuda::MultiHeadAttention::CudaMultiHeadAttentionOp< TensorDataType::BF16 >, Mila::Dnn::Compute::Cuda::MultiHeadAttention::CudaMultiHeadAttentionOp< TensorDataType::FP32 >, Mila::Dnn::Compute::Cuda::RmsNorm::CudaRmsNormOp< TPrecision >, Mila::Dnn::Compute::Cuda::RmsNorm::CudaRmsNormOp< TensorDataType::BF16 >, Mila::Dnn::Compute::Cuda::RmsNorm::CudaRmsNormOp< TensorDataType::FP32 >, Mila::Dnn::Compute::Cuda::Softmax::CudaSoftmaxOp< TPrecision >, Mila::Dnn::Compute::Cuda::Softmax::CudaSoftmaxOp< TensorDataType::BF16 >, Mila::Dnn::Compute::Cuda::Softmax::CudaSoftmaxOp< TensorDataType::FP32 >, Mila::Dnn::Compute::Cuda::Swiglu::CudaSwigluOp< TPrecision >, Mila::Dnn::Compute::Cuda::Swiglu::CudaSwigluOp< TensorDataType::BF16 >, Mila::Dnn::Compute::Cuda::Swiglu::CudaSwigluOp< TensorDataType::FP32 >, Mila::Dnn::Compute::Cuda::TokenEmbedding::CudaTokenEmbeddingOp< TInput, TPrecision >, Mila::Dnn::Compute::Cuda::TokenEmbedding::CudaTokenEmbeddingOp< TensorDataType::INT32, TensorDataType::BF16 >, and Mila::Dnn::Compute::Cuda::TokenEmbedding::CudaTokenEmbeddingOp< TensorDataType::INT32, TensorDataType::FP32 >.

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