Implementation:Microsoft Onnxruntime CUDA GatherElementsGrad
| Knowledge Sources | |
|---|---|
| Domains | Training, CUDA_Kernels |
| Last Updated | 2026-02-10 04:00 GMT |
Overview
Concrete tool for computing the gradient of GatherElements in the ONNX Runtime CUDA training framework.
Description
Implements the GatherElementsGrad operator for CUDA that scatters upstream gradients back to the original data shape. Given the upstream gradient dY (same shape as indices), the original data shape, and indices, it creates a zero-initialized output dX of the original data shape and scatters dY values using GatherElementsGradNonDeterministicImpl. The implementation validates shapes using GatherElements::ValidateInputShapes, coalesces dimensions for efficiency via CoalesceDimensions, and dispatches to type-specific implementations for int32_t and int64_t index types. The operation is inherently non-deterministic due to atomic scatter additions, with a warning logged when deterministic compute is requested. Supports float, double, MLFloat16, and BFloat16. Optionally supports strided index tensors.
Usage
Invoked during the backward pass when the model contains GatherElements operations.
Code Reference
Source Location
- Repository: Microsoft_Onnxruntime
- File: orttraining/orttraining/training_ops/cuda/tensor/gather_elements_grad.cc
- Lines: 1-120
Signature
class GatherElementsGrad : public CudaKernel {
template <typename T> struct ComputeImpl;
Status ComputeInternal(OpKernelContext* context) const;
};
Import
#include "orttraining/training_ops/cuda/tensor/gather_elements_grad.h"
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| dY | Tensor(T) | Yes | Upstream gradient (same shape as indices) |
| shape | Tensor(int64_t) | Yes | Shape of original data tensor (CPU memory) |
| indices | Tensor(Tind) | Yes | Index tensor from forward pass |
Outputs
| Name | Type | Description |
|---|---|---|
| dX | Tensor(T) | Gradient with respect to data (original data shape, zero-initialized then scattered) |
Usage Examples
ONNX_OPERATOR_KERNEL_EX(GatherElementsGrad, kMSDomain, 1, kCudaExecutionProvider,
(*KernelDefBuilder::Create())
.InputMemoryType(OrtMemTypeCPUInput, 1)
.TypeConstraint("T", {DataTypeImpl::GetTensorType<float>(),
DataTypeImpl::GetTensorType<double>(),
DataTypeImpl::GetTensorType<MLFloat16>(),
DataTypeImpl::GetTensorType<BFloat16>()})
.TypeConstraint("I", DataTypeImpl::GetTensorType<int64_t>())
.TypeConstraint("Tind", std::vector<MLDataType>{
DataTypeImpl::GetTensorType<int32_t>(),
DataTypeImpl::GetTensorType<int64_t>()}),
GatherElementsGrad);