KLDiv

English | 简体中文

Product Support Status

Product	Supported
Ascend 950PR&950DT products	√
Atlas A3 products	√
Atlas A2 products	√
Atlas 200I/500 A2 inference products	×
Atlas inference products	√
Atlas training products	√

Function Description

• Operator function: computes the KL divergence loss (Kullback-Leibler Divergence).

• Shell operator positioning: KLDiv is a shell routing operator. Its registered name KLDiv exactly matches the op type of the legacy GE graph and is used for compatibility with legacy graph migration. If you have new computation requirements, use KLDivV2 instead. The computation kernel of this operator purely forwards the log_target=False DAG of KLDivV2 and always computes with the log_target=False semantics.

• Computation formula (always log_target=False):

$$los{s}_i=\{\begin{array}{ll}0& \text{if }targe{t}_i=0\\ NaN& \text{if }targe{t}_i<0\\ targe{t}_i\times \left(\log (targe{t}_i)-x_i\right)& \text{else}\end{array}$$

• The output y is determined by reduction:

$$y=\{\begin{array}{ll}loss& \text{if }reduction=\text{"none"}\\ {\sum }_{i}los{s}_i& \text{if }reduction=\text{"sum"}\\ \frac{{\sum }_{i}los{s}_i}{numel}& \text{if }reduction=\text{"mean"}\\ \frac{{\sum }_{i}los{s}_i}{di{m}_0}& \text{if }reduction=\text{"batchmean"}\end{array}$$

• When reduction is "none", y has the same shape as x; when reduction is "mean", "sum", or "batchmean", y is a scalar (with an empty shape).

Parameter Description

Parameter Name	Input/Output/Attribute	Description	Data Type	Data Format
x	Input	Input tensor x (log probabilities) in the formula.	FLOAT, FLOAT16, BFLOAT16	ND
target	Input	Input tensor target (probabilities) in the formula, which has the same shape and dtype as x.	FLOAT, FLOAT16, BFLOAT16	ND
reduction	Required attribute	Specifies the reduction mode of the output. Only none, mean, sum, and batchmean are supported. There is no default value, and it must be explicitly set when building a graph.	STRING	-
y	Output	Output tensor y in the formula. When reduction is "none", it has the same shape as x; for other values, it is a scalar.	FLOAT, FLOAT16, BFLOAT16	ND

• The operator prototype does not declare the log_target attribute (the v1 prototype does not have this attribute either), and the operator always computes with the log_target=False semantics. The log_target attribute carried in a legacy graph is ignored as an unknown attribute.

Constraint Description

• x and target must have the same shape and dtype.
• reduction is a required attribute, and only none, mean, sum, and batchmean are supported.
• The prototype declares DT_DOUBLE support on paper, but the dtype whitelist in the tiling phase rejects double. DT_DOUBLE is actually not supported.
• The output is 0 at positions where target equals 0, and NaN at positions where target is less than 0 (value range constraint, which is not verified by tiling/infershape; callers must ensure the input semantics are valid).
• FLOAT16 and BFLOAT16 inputs are promoted to FLOAT32 for computation inside the kernel, and the results are cast back to the original dtype by round-to-nearest-even.
• Output shape: when reduction is "none", y has the same shape as x; when reduction is "mean", "sum", or "batchmean", y is a scalar.
• This operator is a shell routing operator and has no aclnn API: legacy graph migration uses the atc graph compilation path. For PyTorch adaptation, use the aclnnKlDiv API in the kl_div_v2 directory of ops-math.

Invocation Description

Invocation Method	Invocation Sample	Description
Graph-mode invocation	test_geir_kl_div	Refer to operator invocation to complete the operator compilation and verification.

• The graph-mode example invokes this operator in the op::KLDiv graph-building mode through the operator IR.