Skip to content

fix(likelihoods): pass targets through to DirichletClassificationLikelihood.get_fantasy_likelihood - #2758

Merged
gpleiss merged 8 commits into
cornellius-gp:mainfrom
umi008:fix-2579-dirichlet-fantasy-likelihood
Jul 10, 2026
Merged

fix(likelihoods): pass targets through to DirichletClassificationLikelihood.get_fantasy_likelihood#2758
gpleiss merged 8 commits into
cornellius-gp:mainfrom
umi008:fix-2579-dirichlet-fantasy-likelihood

Conversation

@umi008

@umi008 umi008 commented Jul 9, 2026

Copy link
Copy Markdown
Contributor

Closes #2579

Root cause

DirichletClassificationLikelihood.get_fantasy_likelihood requires a targets kwarg, but ExactGP.get_fantasy_model never passes targets into fantasy_kwargs. Additionally, the method reads the wrong kwarg key ("noise" instead of "targets" — copy-paste from the parent FixedNoiseGaussianLikelihood), and the error message names the wrong class.

Fixes

  1. gpytorch/models/exact_gp.py: add targets to fantasy_kwargs before calling get_fantasy_likelihood. targets is popped before being forwarded to get_fantasy_strategy, which receives them positionally.
  2. gpytorch/likelihoods/gaussian_likelihood.py:
    • Error message class name corrected: FixedNoiseGaussianLikelihood.fantasize -> DirichletClassificationLikelihood.get_fantasy_likelihood.
    • kwargs.get("noise") -> kwargs.get("targets").
    • When called from get_fantasy_model, targets arrives as the transformed 2D one-hot tensor (not the raw 1D class labels). Apply argmax(dim=0) to recover the raw labels before _prepare_targets.

Changes

File Change
gpytorch/models/exact_gp.py Pass targets through to get_fantasy_likelihood; pop before forwarding to get_fantasy_strategy
gpytorch/likelihoods/gaussian_likelihood.py Three fixes: error class name, kwarg key, transformed-targets recovery
test/likelihoods/test_gaussian_likelihood.py 4 regression tests, one marked expectedFailure for a separate get_fantasy_strategy view-size bug (tracked separately)

Test plan

  • 4 new regression tests pass (one expectedFailure for the downstream view-size bug)
  • Full test/likelihoods suite passes (35 tests)
  • pre-commit run --files gpytorch/models/exact_gp.py gpytorch/likelihoods/gaussian_likelihood.py test/likelihoods/test_gaussian_likelihood.py passes

umi008 added 6 commits July 9, 2026 12:51
…antasy model)

Reproduces the RuntimeError when calling get_fantasy_model on a model
with DirichletClassificationLikelihood. The test proves the bug exists
before applying the fix.

Co-Authored-By: none
…elihood

The error message incorrectly named FixedNoiseGaussianLikelihood.fantasize
when the actual failing method is DirichletClassificationLikelihood.
get_fantasy_likelihood.

Co-Authored-By: none
…asy_likelihood

Line 450 was reading kwargs.get("noise") which is a copy-paste error
from the parent class. DirichletClassificationLikelihood.get_fantasy_likelihood
needs to read kwargs.get("targets") to get the raw class labels for
_prepare_targets.

Co-Authored-By: none
…onLikelihood.get_fantasy_likelihood

When get_fantasy_likelihood is called via get_fantasy_model, the targets
passed in are the already-transformed targets (2D, one-hot per class)
from the model, but _prepare_targets expects raw class labels (1D).

Convert 2D targets to class indices via argmax before calling
_prepare_targets. This unblocks the fantasy model flow that was
broken at multiple points (#2579).
…ntasy_model

DirichletClassificationLikelihood.get_fantasy_likelihood requires the
'targets' kwarg to reconstruct the fantasy noise from class labels.
Add it to fantasy_kwargs alongside 'noise' so the call succeeds (#2579).

Pop 'targets' from fantasy_kwargs before passing to
get_fantasy_strategy because that call already receives targets
positionally.
…sting issue

Mark the end-to-end get_fantasy_model test as expectedFailure: it
exercises a separate, pre-existing bug in
DefaultPredictionStrategy.get_fantasy_strategy (view size mismatch on
batched models) that is outside the scope of #2579.

Add a new direct test test_dirichlet_get_fantasy_likelihood_accepts_targets
that exercises the exact code path fixed by this PR (get_fantasy_likelihood
accepts 'targets' kwarg and returns a DirichletClassificationLikelihood).

Also applies ufmt reformatting from pre-commit.

@gpleiss gpleiss left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

See review, and fix linting.

Comment thread gpytorch/likelihoods/gaussian_likelihood.py Outdated
Comment thread test/likelihoods/test_gaussian_likelihood.py
Address review feedback from gpleiss on PR #2758:

- gaussian_likelihood.py: Change argmax(dim=0) to argmax(dim=-2) for
  batch compatibility — negative index always targets the class dimension
  regardless of leading batch dimensions.
- test: Rewrite test_dirichlet_fantasy_model_creation to use batched
  train_x/train_y (batch_size=2) to ensure indexing operations are
  batch compatible.
- test: Remove unused fant_labels variable (F841 lint fix).
@umi008
umi008 requested a review from gpleiss July 10, 2026 21:07
@gpleiss
gpleiss merged commit 6272eda into cornellius-gp:main Jul 10, 2026
7 checks passed
@gpleiss

gpleiss commented Jul 10, 2026

Copy link
Copy Markdown
Member

Thanks @umi008

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug] get_fantasy_likelihood method broken for DirichletClassificationLikelihood

2 participants