Medium Primitives

Implement KL Divergence

KL divergence compares two probability distributions in a specified direction. For D_KL(P || Q), weight each log-ratio by P, then sum over outcomes:

D_KL(P || Q) = sum(p[i] * log(p[i] / q[i]))

Use the natural logarithm, so the result is measured in nats. Reversing P and Q usually changes the result; KL divergence is not a symmetric distance.

Your task

Implement kl_divergence(p, q) from elementwise arithmetic, logarithms and a sum. Do not use a built-in KL-divergence helper.

The inputs are one-dimensional floating-point tensors of the same nonempty shape, dtype and device. Each contains strictly positive probabilities that sum to one. Inputs with zero or negative probabilities are outside this task.

This exercise uses a fixed additive epsilon of 1e-10. Apply it exactly as follows, including in the outer multiplier:

p_safe = p + 1e-10
q_safe = q + 1e-10
result = sum(p_safe * log(p_safe / q_safe))

Return a scalar tensor with shape (), preserving the input dtype and leaving p and q unchanged. Do not average across outcomes, swap the inputs, renormalize after adding epsilon, or clamp the final result to zero.

For p=[0.8, 0.2] and q=[0.5, 0.5], the result is approximately 0.192745. The first outcome contributes positively and the second contributes negatively; sum both terms. Equal input distributions give zero.

What does epsilon change?

The exact KL formula is the expression at the top. The required computation instead uses slightly shifted values whose sums exceed one. It approximates KL when epsilon is small relative to the probabilities, but can differ noticeably when a probability is close to epsilon. Use the specified shifted formula for grading, not a different clipping or smoothing rule.

Exact KL is nonnegative for normalized distributions. Floating-point sums can have small rounding errors near zero, so do not use a final clamp to conceal them. This exercise does not define the infinite-divergence case where Q is zero at an outcome with positive P.

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Compare two probability distributions

Kullback–Leibler (KL) divergence compares a distribution P with a distribution Q over the same outcomes:

$$D_{KL}(P\|Q)=\sum_i P_i\log\frac{P_i}{Q_i}.$$

Each term compares the log probabilities assigned to one outcome, weighted by how likely that outcome is under P. Using natural logarithms measures the result in nats.

For P = [1, 0] and Q = [0.5, 0.5], only the first outcome has weight, giving log(1 / 0.5) = log(2). If both distributions are identical, the divergence is zero.

The order matters

Swapping P and Q changes both the weights and the probability ratios. KL is therefore generally asymmetric, so it is not a distance metric. In the example above, reversing the arguments gives infinite divergence: Q assigns positive probability to an outcome that P says is impossible.

Handle zeros explicitly

A term with P_i = 0 contributes zero under the standard limiting convention. Computing it directly as 0 * log(0) can instead produce NaN. If P_i > 0 and Q_i = 0, the mathematical divergence is infinite.

Adding epsilon or clamping probabilities changes these boundary cases. Follow the exercise's convention rather than silently replacing infinite values.

Check the representation and reduction

Probabilities, log probabilities, and logits need different calculations. Convert or use the appropriate formula; logits cannot be inserted directly into the probability formula above.

For a batch, sum over the outcomes of each distribution first, then apply the requested reduction over examples. In exact arithmetic, KL is nonnegative for normalized distributions. Floating-point rounding can produce tiny negative values near zero. A substantially negative result suggests an error in signs, normalization, or reduction.

Detecting runtime…