Fix sparse categorical cross entropy for softmax outputs

[dajiaohuang]

Axon.Losses.categorical_cross_entropy/3 ignores the :sparse option when y_pred carries cached logits metadata from Axon.Activations.softmax/1. That metadata only exists inside a compiled defn, so a softmax-activated model trained with sparse integer targets silently takes the dense code path and produces a wrong loss:

defmodule Repro do
  import Nx.Defn

  defn loss(y_true, logits) do
    y_pred = Axon.Activations.softmax(logits)
    Axon.Losses.categorical_cross_entropy(y_true, y_pred, sparse: true, reduction: :mean)
  end
end

y_true = Nx.tensor([1, 0, 2])
logits = Nx.tensor([[15.0, 2.0, -22.0], [-5.0, -2.0, 3.0], [2.0, 1.96, 1.20]])

Repro.loss(y_true, logits)
#=> 28.753393   (expected 7.562242)

The same computation with from_logits: true returns the correct 7.562242, because that branch forwards :sparse to softmax_cross_entropy_from_logits/3. The logits metadata branch a few lines below does not:

from_logits: true -> softmax_cross_entropy_from_logits(y_true, logits, sparse: sparse)
logits?           -> softmax_cross_entropy_from_logits(y_true, logits)  # :sparse dropped

Since softmax_cross_entropy_from_logits/3 defaults to sparse: false, the sparse targets are then treated as dense one-hot labels, which is what makes the result wrong rather than raising.

This forwards :sparse in the metadata branch as well. The default is unchanged: :sparse defaults to false in both categorical_cross_entropy/3 and softmax_cross_entropy_from_logits/3, so dense callers see identical behavior.

The existing sparse tests in test/axon/losses_test.exs call Axon.Activations.softmax/1 eagerly at test scope, where the metadata is lost and the correct branch is taken — which is why this went unnoticed. The regression test added here calls softmax inside a compiled defn so the metadata survives, and checks the sparse and dense forms against the existing 7.562242 reference value.

Validation:

• mix test test/axon/losses_test.exs — 39 doctests, 18 tests, 0 failures
• mix test — 200 doctests, 738 tests, 4 failures; the 4 failures are pre-existing in this checkout (one gelu doctest float-precision difference and three AxonTest inspect assertions affected by a CRLF working tree) and are unchanged from before the patch
• mix format --check-formatted on both changed files — clean

Reverting the one-line change makes the new test fail, so it does exercise the reported path.

This is the idiom Axon.Loop's own documentation suggests for integer class labels (Axon.Loop.trainer/4 and the shape-mismatch hint in lib/axon/loop.ex), and it is reached through the normal Axon.Loop.train_step/4 path, so a softmax-activated classifier trained on sparse labels silently reports an inflated loss and trains on the wrong gradients.