Skip to content

Remove the dead cache_trainset_representation flag - #1250

Open
jmkuebler wants to merge 3 commits into
mainfrom
jonas/drop_cache_trainset_representation
Open

Remove the dead cache_trainset_representation flag#1250
jmkuebler wants to merge 3 commits into
mainfrom
jonas/drop_cache_trainset_representation

Conversation

@jmkuebler

@jmkuebler jmkuebler commented Sep 8, 2026

Copy link
Copy Markdown
Contributor

Issue

RES-2743 (base of #1219).

Motivation and Context

cache_trainset_representation has been dead since InferenceEngineCacheKV was removed in #1057. Every architecture's get_architecture accepts it and immediately dels it; the KV cache is owned by the inference engine and selected at call time through forward(). The flag survived only because the ArchitectureModule protocol still declared it, which forced every implementation, every loader, and every test call site to carry it.

This removes it end to end:

  • ArchitectureModule.get_architecture(config) no longer takes the kwarg; the four architectures drop the parameter and their "accepted for interface compatibility but ignored" notes.
  • load_model_criterion_config, load_model and _build_model lose the parameter.
  • The opt-in built-model cache (TABPFN_MODEL_CACHE_SIZE) no longer gates on the flag. The gate assumed a fit_with_cache build mutates the module, which it does not, so that fit mode now uses the cache like every other one.
  • initialize_tabpfn_model loses its fit_mode argument, whose only job was computing the flag.
  • Tests: every call site, the three fake ArchitectureModules, the "mutating build is never cached" test, and the load_model signature tripwire.

#1219 is stacked on top and shrinks to keying the cache on device and dtype.

fomo-fitting side: PriorLabs/fomo-fitting#3776 drops the parameter from every architecture there. It is a draft until this merges and the tabpfn pin is bumped.


Public API Changes

  • No Public API changes

  • Yes, Public API changes (Details below)

  • ArchitectureModule.get_architecture no longer accepts cache_trainset_representation. External architectures implementing the protocol must drop the parameter.

  • tabpfn.model_loading.load_model and load_model_criterion_config no longer accept cache_trainset_representation.

  • tabpfn.base.initialize_tabpfn_model no longer accepts fit_mode.

  • With TABPFN_MODEL_CACHE_SIZE set, fit_mode="fit_with_cache" now uses the built-model cache. The default (cache off) is unchanged.


How Has This Been Tested?

  • ruff check / format clean.
  • Cache off (shipping default): tests/test_classifier_interface.py and tests/test_regressor_interface.py pass on cpu and cuda; the touched unit files (test_model_cache, test_model_loading, test_config, test_inference_config, test_inference, test_architectures/, test_save_load_fitted_model) pass.
  • Cache on, the interface suites fail with dtype/device collisions, as they did before this change (see fix: key the built-model cache on device and inference dtype #1219, which fixes that and is stacked on this PR).

Checklist

  • The changes have been tested locally.
  • Documentation has been updated (if the public API or usage changes).
  • A changelog entry has been added (see changelog/README.md), or "no changelog needed" label requested.
  • The code follows the project's style guidelines.
  • I have considered the impact of these changes on the public API.

🤖 Generated with Claude Code

jmkuebler and others added 2 commits September 8, 2026 09:07
Every architecture's get_architecture accepted the flag and deleted it
unused: model-internal caching went away with InferenceEngineCacheKV
(#1057), and the KV cache is owned by the inference engine. Drop it from
the ArchitectureModule protocol, the loaders, and the tests.

The built-model cache (TABPFN_MODEL_CACHE_SIZE) no longer gates on the
flag, so fit_mode="fit_with_cache" uses it like every other mode.
initialize_tabpfn_model loses its fit_mode argument, whose only job was
computing the flag.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
jmkuebler added a commit that referenced this pull request Sep 8, 2026
…d by #1250

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
@jmkuebler
jmkuebler requested a review from oscarkey September 8, 2026 09:48
@jmkuebler

Copy link
Copy Markdown
Contributor Author

cc @priorphil as you authored #1057 any reason to keep cache_trainset_representation?

oscarkey
oscarkey previously approved these changes Sep 8, 2026

@oscarkey oscarkey left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

sweeeet

Without the bare *, a stale positional call that still passed fit_mode as
the third argument would silently bind it to softmax_temperature_override,
and with several checkpoints that non-None value would also suppress the
temperature-disagreement check. Now it fails at the call.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
jmkuebler added a commit that referenced this pull request Sep 8, 2026
…d by #1250

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Reapproving after a push, as this is a PR from an internal author and was already approved.

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.

2 participants