LightPFN 1.0.0: public model card, report, provenance; remove the private 0.1.0 package files

LICENSE

@@ -1,202 +1,202 @@
-
-                                 Apache License
-                           Version 2.0, January 2004
-                        http://www.apache.org/licenses/
-
-   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
-
-   1. Definitions.
-
-      "License" shall mean the terms and conditions for use, reproduction,
-      and distribution as defined by Sections 1 through 9 of this document.
-
-      "Licensor" shall mean the copyright owner or entity authorized by
-      the copyright owner that is granting the License.
-
-      "Legal Entity" shall mean the union of the acting entity and all
-      other entities that control, are controlled by, or are under common
-      control with that entity. For the purposes of this definition,
-      "control" means (i) the power, direct or indirect, to cause the
-      direction or management of such entity, whether by contract or
-      otherwise, or (ii) ownership of fifty percent (50%) or more of the
-      outstanding shares, or (iii) beneficial ownership of such entity.
-
-      "You" (or "Your") shall mean an individual or Legal Entity
-      exercising permissions granted by this License.
-
-      "Source" form shall mean the preferred form for making modifications,
-      including but not limited to software source code, documentation
-      source, and configuration files.
-
-      "Object" form shall mean any form resulting from mechanical
-      transformation or translation of a Source form, including but
-      not limited to compiled object code, generated documentation,
-      and conversions to other media types.
-
-      "Work" shall mean the work of authorship, whether in Source or
-      Object form, made available under the License, as indicated by a
-      copyright notice that is included in or attached to the work
-      (an example is provided in the Appendix below).
-
-      "Derivative Works" shall mean any work, whether in Source or Object
-      form, that is based on (or derived from) the Work and for which the
-      editorial revisions, annotations, elaborations, or other modifications
-      represent, as a whole, an original work of authorship. For the purposes
-      of this License, Derivative Works shall not include works that remain
-      separable from, or merely link (or bind by name) to the interfaces of,
-      the Work and Derivative Works thereof.
-
-      "Contribution" shall mean any work of authorship, including
-      the original version of the Work and any modifications or additions
-      to that Work or Derivative Works thereof, that is intentionally
-      submitted to Licensor for inclusion in the Work by the copyright owner
-      or by an individual or Legal Entity authorized to submit on behalf of
-      the copyright owner. For the purposes of this definition, "submitted"
-      means any form of electronic, verbal, or written communication sent
-      to the Licensor or its representatives, including but not limited to
-      communication on electronic mailing lists, source code control systems,
-      and issue tracking systems that are managed by, or on behalf of, the
-      Licensor for the purpose of discussing and improving the Work, but
-      excluding communication that is conspicuously marked or otherwise
-      designated in writing by the copyright owner as "Not a Contribution."
-
-      "Contributor" shall mean Licensor and any individual or Legal Entity
-      on behalf of whom a Contribution has been received by Licensor and
-      subsequently incorporated within the Work.
-
-   2. Grant of Copyright License. Subject to the terms and conditions of
-      this License, each Contributor hereby grants to You a perpetual,
-      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
-      copyright license to reproduce, prepare Derivative Works of,
-      publicly display, publicly perform, sublicense, and distribute the
-      Work and such Derivative Works in Source or Object form.
-
-   3. Grant of Patent License. Subject to the terms and conditions of
-      this License, each Contributor hereby grants to You a perpetual,
-      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
-      (except as stated in this section) patent license to make, have made,
-      use, offer to sell, sell, import, and otherwise transfer the Work,
-      where such license applies only to those patent claims licensable
-      by such Contributor that are necessarily infringed by their
-      Contribution(s) alone or by combination of their Contribution(s)
-      with the Work to which such Contribution(s) was submitted. If You
-      institute patent litigation against any entity (including a
-      cross-claim or counterclaim in a lawsuit) alleging that the Work
-      or a Contribution incorporated within the Work constitutes direct
-      or contributory patent infringement, then any patent licenses
-      granted to You under this License for that Work shall terminate
-      as of the date such litigation is filed.
-
-   4. Redistribution. You may reproduce and distribute copies of the
-      Work or Derivative Works thereof in any medium, with or without
-      modifications, and in Source or Object form, provided that You
-      meet the following conditions:
-
-      (a) You must give any other recipients of the Work or
-          Derivative Works a copy of this License; and
-
-      (b) You must cause any modified files to carry prominent notices
-          stating that You changed the files; and
-
-      (c) You must retain, in the Source form of any Derivative Works
-          that You distribute, all copyright, patent, trademark, and
-          attribution notices from the Source form of the Work,
-          excluding those notices that do not pertain to any part of
-          the Derivative Works; and
-
-      (d) If the Work includes a "NOTICE" text file as part of its
-          distribution, then any Derivative Works that You distribute must
-          include a readable copy of the attribution notices contained
-          within such NOTICE file, excluding those notices that do not
-          pertain to any part of the Derivative Works, in at least one
-          of the following places: within a NOTICE text file distributed
-          as part of the Derivative Works; within the Source form or
-          documentation, if provided along with the Derivative Works; or,
-          within a display generated by the Derivative Works, if and
-          wherever such third-party notices normally appear. The contents
-          of the NOTICE file are for informational purposes only and
-          do not modify the License. You may add Your own attribution
-          notices within Derivative Works that You distribute, alongside
-          or as an addendum to the NOTICE text from the Work, provided
-          that such additional attribution notices cannot be construed
-          as modifying the License.
-
-      You may add Your own copyright statement to Your modifications and
-      may provide additional or different license terms and conditions
-      for use, reproduction, or distribution of Your modifications, or
-      for any such Derivative Works as a whole, provided Your use,
-      reproduction, and distribution of the Work otherwise complies with
-      the conditions stated in this License.
-
-   5. Submission of Contributions. Unless You explicitly state otherwise,
-      any Contribution intentionally submitted for inclusion in the Work
-      by You to the Licensor shall be under the terms and conditions of
-      this License, without any additional terms or conditions.
-      Notwithstanding the above, nothing herein shall supersede or modify
-      the terms of any separate license agreement you may have executed
-      with Licensor regarding such Contributions.
-
-   6. Trademarks. This License does not grant permission to use the trade
-      names, trademarks, service marks, or product names of the Licensor,
-      except as required for reasonable and customary use in describing the
-      origin of the Work and reproducing the content of the NOTICE file.
-
-   7. Disclaimer of Warranty. Unless required by applicable law or
-      agreed to in writing, Licensor provides the Work (and each
-      Contributor provides its Contributions) on an "AS IS" BASIS,
-      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
-      implied, including, without limitation, any warranties or conditions
-      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
-      PARTICULAR PURPOSE. You are solely responsible for determining the
-      appropriateness of using or redistributing the Work and assume any
-      risks associated with Your exercise of permissions under this License.
-
-   8. Limitation of Liability. In no event and under no legal theory,
-      whether in tort (including negligence), contract, or otherwise,
-      unless required by applicable law (such as deliberate and grossly
-      negligent acts) or agreed to in writing, shall any Contributor be
-      liable to You for damages, including any direct, indirect, special,
-      incidental, or consequential damages of any character arising as a
-      result of this License or out of the use or inability to use the
-      Work (including but not limited to damages for loss of goodwill,
-      work stoppage, computer failure or malfunction, or any and all
-      other commercial damages or losses), even if such Contributor
-      has been advised of the possibility of such damages.
-
-   9. Accepting Warranty or Additional Liability. While redistributing
-      the Work or Derivative Works thereof, You may choose to offer,
-      and charge a fee for, acceptance of support, warranty, indemnity,
-      or other liability obligations and/or rights consistent with this
-      License. However, in accepting such obligations, You may act only
-      on Your own behalf and on Your sole responsibility, not on behalf
-      of any other Contributor, and only if You agree to indemnify,
-      defend, and hold each Contributor harmless for any liability
-      incurred by, or claims asserted against, such Contributor by reason
-      of your accepting any such warranty or additional liability.
-
-   END OF TERMS AND CONDITIONS
-
-   APPENDIX: How to apply the Apache License to your work.
-
-      To apply the Apache License to your work, attach the following
-      boilerplate notice, with the fields enclosed by brackets "[]"
-      replaced with your own identifying information. (Don't include
-      the brackets!)  The text should be enclosed in the appropriate
-      comment syntax for the file format. We also recommend that a
-      file or class name and description of purpose be included on the
-      same "printed page" as the copyright notice for easier
-      identification within third-party archives.
-
-   Copyright [yyyy] [name of copyright owner]
-
-   Licensed under the Apache License, Version 2.0 (the "License");
-   you may not use this file except in compliance with the License.
-   You may obtain a copy of the License at
-
-       http://www.apache.org/licenses/LICENSE-2.0
-
-   Unless required by applicable law or agreed to in writing, software
-   distributed under the License is distributed on an "AS IS" BASIS,
-   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-   See the License for the specific language governing permissions and
-   limitations under the License.
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

LightPFN_report.pdf

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8bea6f04f76fa4d082d8884445169b20c9e55b901543c205de46d9fa7345d7b0
-size 460153
+oid sha256:274be6425d074c984ce49cee56f40a0801421017a2851febe916ac7bbc4718f4
+size 331223

NOTICE

@@ -1,6 +1,11 @@
 LightPFN
-Copyright 2026 Giorgio (GioOtto)
+Copyright 2026 Giorgio Ottoboni
 
-Code and released model weights are licensed under the Apache License, Version 2.0.
-The model was trained from scratch using synthetic data, without distillation
-or training on weights or outputs from other tabular foundation models.
+This product includes software and model weights licensed under the Apache License, Version 2.0.
+
+The model was trained from scratch on synthetic data only, without distillation and without weights or
+outputs of other tabular foundation models.
+
+The optional training code (lightpfn.prior) uses the structural causal model prior of TabICL
+(https://github.com/soda-inria/tabicl, BSD-3-Clause, Copyright (c) 2025, Soda team @ Inria) as an installed
+dependency; TabICL code is not redistributed in this repository or in the published wheel.

README.md

@@ -4,122 +4,96 @@ library_name: lightpfn
 pipeline_tag: tabular-classification
 tags:
   - tabular
-  - classification
+  - tabular-classification
   - in-context-learning
+  - prior-data-fitted-network
+  - foundation-model
   - synthetic-data
-  - pytorch
+  - scikit-learn
+  - vulkan
   - safetensors
 ---
 
-# LightPFN 0.1.0
+# LightPFN
 
-**4,603,088 parameters**, pretrained from scratch **only on synthetic data**.
-This repository contains the final `final_long` model, safe inference weights,
-and the installable Python package. Both **code and weights are Apache 2.0**.
-The initial repository is private; it is not yet a public release or a PyPI upload.
+LightPFN is a small tabular foundation model for classification: a 4,603,088-parameter in-context learner
+pretrained only on synthetic data. `fit` stores the training set as context and `predict_proba` answers in one
+forward pass, with no training on your data and no hyperparameters to tune. It runs on CPU, CUDA, ROCm and, through
+its own Vulkan kernels, on AMD, Intel and NVIDIA GPUs.
 
-Preprint: [A Sling Against Giants: LightPFN, a 4.6M-parameter tabular in-context
-classifier designed to stay small](./LightPFN_report.pdf), Giorgio Ottoboni,
-5 October 2026.
+- Code, documentation and training pipeline: [github.com/GioOtto/LightPFN](https://github.com/GioOtto/LightPFN)
+- Package: [pypi.org/project/LightPFN](https://pypi.org/project/LightPFN/)
+- Technical report: [LightPFN_report.pdf](./LightPFN_report.pdf)
+- License: Apache 2.0, code and weights
 
-No distillation was used. No weights or outputs from other tabular foundation
-models were used in training. External tabular models were evaluation baselines
-only. The release excludes their code/weights, synthetic priors and training
-dependencies.
-
-## Installation
-
-Authenticate with an account authorized to read this private repository:
+## Usage
 
 ```bash
-hf auth login
-hf download ueuegio/LightPFN dist/lightpfn-0.1.0-py3-none-any.whl --local-dir .
-pip install "./dist/lightpfn-0.1.0-py3-none-any.whl[sklearn,hf]"
+pip install LightPFN
 ```
 
-Install `huggingface_hub` first if the `hf` command is unavailable. For CPU-only
-use, install torch from PyTorch's official CPU wheel index before the package.
-An optional Vulkan extra provides AMD/Intel/NVIDIA GPU inference without CUDA:
-`pip install "./dist/lightpfn-0.1.0-py3-none-any.whl[sklearn,hf,vulkan]"`.
-
 ```python
 from lightpfn import LightPFNClassifier
 
-clf = LightPFNClassifier(device="auto", n_estimators=4, random_state=0)
-clf.fit(X_train, y_train)
+clf = LightPFNClassifier(n_estimators=4, random_state=0)
+clf.fit(X_train, y_train)          # downloads these weights at a pinned commit on the first fit
 proba = clf.predict_proba(X_test)
-pred = clf.predict(X_test)
 ```
 
-The package pins the weights/config to an immutable commit. The first fit downloads
-them using the HF cache. `local_files_only=True` supports an already cached offline
-model. Construction and sklearn cloning do no downloads or GPU initialization.
-Inputs are dense numeric tables; missing values can be NaN. Encode string
-categories before fitting. Native categorical handling is not implemented.
-The deprecated `cat` mask does not change the computation and warns when used
-to flag categories.
-
-The base package needs only torch and numpy; sklearn, Hub/safetensors and Vulkan
-are optional extras. The inference package imports no prior or external foundation
-model. The root `model.safetensors` stores tensors and `config.json` stores the
-architecture. Legacy tensor/dict `.pt` files load only with `weights_only=True`.
-No unrestricted pickle fallback or remote-code loading is provided.
-
-## Training
-
-Architecture: B4, summary-token rows, three row-refinement rounds, seven ICL
-blocks, retrieval decoder. The synthetic mix is graph v3 90% / rule v1 10%,
-without augmentation. Model size was fixed below 5M throughout development.
+`X` can be a NumPy array or a pandas DataFrame with numeric, categorical, string and boolean columns and missing
+values. `device="auto"` uses a CUDA or ROCm GPU, else a Vulkan GPU (`pip install "LightPFN[vulkan]"`), else the CPU.
+The [user guide](https://github.com/GioOtto/LightPFN/blob/main/docs/en/GUIDE.md) covers every option.
 
-- Base stage: 120,000 steps, 64 tasks/step, approximately 7.7M tasks seen, on two
-  RTX 5090 GPUs; tables of at most 2,048 rows.
-- Long-context stage: 39,250 steps, LR 1e-4 to 1e-5, tables up to 60,000 rows.
-- Released weights: EMA `final_long/ema_step039250.pt`, converted losslessly to
-  FP32 safetensors. Training optimizer/state is excluded.
+## Model
 
-`provenance.json` records source checkpoint hash, tensor artifact hashes,
-architecture size, software versions, source-code revision and preprint hash.
+| | |
+|---|---|
+| Task | classification, 2 to 10 classes |
+| Parameters | 4,603,088 (18 MB, float32) |
+| Architecture | cell embedding (value, rank, missing flag), two induced column stages, row refinement with four summary tokens, row compression, seven in-context blocks, retrieval decoder |
+| Pretraining data | synthetic only: 90% structural causal graph prior, 10% rule prior (XOR, parity, lookup, trees); 7.68 million task draws from 4.03 million distinct tasks |
+| Training | 120,000 steps of 64 tasks on two RTX 5090 (tables up to 2,048 rows), then 39,250 steps on tables up to 60,000 rows |
+| Not used in training | real datasets, distillation, weights or outputs of other tabular foundation models |
 
 ## Evaluation
 
-Variants were selected on held-out synthetic tasks (D1), interaction probes (D2)
-and 55 OpenML-CC18 tasks outside TabArena (D3). TabArena was inspected only for
-the finalists. These are local experiments, **not an official TabArena submission**.
-GBDT comparisons use library defaults, not tuned TabArena configurations.
+The official TabArena-Lite evaluation includes default, tuned and ensembled baselines. Other evaluations
+use default baselines; AUC differences have 95% paired bootstrap intervals.
 
-| Evaluation | Result |
+| Benchmark | Result |
 |---|---|
-| D3, 55 tasks, <=1,000 training rows, one estimator | Mean AUC 0.9109 |
-| D3 vs default CatBoost | +0.86 AUC points, paired 95% CI [0.42, 1.40] |
-| TabArena classification, 38 tasks, official splits, first repeat, four estimators | Mean AUC 0.8581, mean rank 2.50 |
-| Same TabArena run vs default CatBoost | Lower task error on 76% of tasks |
+| Official TabArena-Lite pipeline, 38 classification datasets, default configuration, four estimators | Elo 1420 (+67 / -66), 25th of 99 methods, 38 successful tasks, none imputed; above GBDT point estimates, intervals overlap tuned CatBoost; author-run, pending maintainer full-benchmark verification |
+| 55 OpenML-CC18 datasets outside TabArena (at most 1,000 rows, one estimator) | mean AUC 0.911, +0.86 [0.42, 1.40] over CatBoost, +1.8 to +2.1 over LightGBM, XGBoost and random forest |
+| TabArena, 38 classification tasks, official splits, first repeat (our harness), four estimators | mean AUC 0.858, rank 2.50 of 7, behind TabICLv2 (0.864), lower error than CatBoost on 76% of tasks |
+| 13 OpenML datasets of 50,000 to 2.2M rows, 10,000 to 100,000 training rows | mean AUC difference from CatBoost between -0.34 and +0.28 points, intervals include zero |
 
-Task error is 1-AUC for binary tasks and log loss for multiclass tasks. On this
-full TabArena evaluation, TabICLv2 ranked first (AUC 0.8637, mean rank 1.58).
-TabArena-lite and full results should not be conflated. The full report and
-evaluation outputs are in the source repository.
+Full tables, timings and the official TabArena results:
+[docs/en/RESULTS.md](https://github.com/GioOtto/LightPFN/blob/main/docs/en/RESULTS.md).
 
 ## Limitations
 
-Classification for **2-10 classes**, without regression. A single observed class
-returns its constant probability. The default context limit is 20,000 rows;
-larger datasets use a stratified subset for each ensemble member. Large contexts
-cost substantially more memory/time. Categorical-heavy datasets are a known
-weakness: in_vehicle_coupon_recommendation and Amazon_employee_access trailed
-default CatBoost by about 4 AUC points in the full evaluation. CPU inference is
-sometimes slower than default GBDTs (median 15/54 seconds on medium/large
-TabArena splits versus CatBoost's 4/7 seconds on the evaluated server).
-
-Inference folding, cell blocking and estimator batching preserve predictions
-within FP32 tolerances. Row permutation or query batch changes may shift
-probabilities around 1e-7; two strict sklearn invariance checks document this
-tolerance. GPU adapters have finite storage-buffer limits; device="auto" can
-fall back to CPU when a Vulkan context exceeds those limits.
-
-## Source and license
-
-Source: [GioOtto/gioPFN](https://github.com/GioOtto/gioPFN).
-The repo's code may require separate access while it is private.
-Code and released weights: [Apache License 2.0](./LICENSE), with [NOTICE](./NOTICE).
-Third-party dependencies retain their own licenses, listed in
-`package/docs/DEPENDENCIES.md`.
+- Classification only (2 to 10 classes); regression is planned for version 2.
+- Categorical columns are read as ordinal codes. On tables dominated by high-cardinality categorical columns
+  CatBoost is ahead; native categorical handling is planned for version 2.
+- Above 20,000 training rows each estimator reads a stratified subsample (`max_context`).
+- CPU time grows with the context length; on large tables a GPU is much faster.
+
+## Files
+
+| File | Content |
+|---|---|
+| `model.safetensors`, `config.json` | weights (float32) and architecture of the released model |
+| `provenance.json` | hashes of the weights and of the source checkpoint |
+| `LightPFN_report.pdf` | technical report |
+| `LICENSE`, `NOTICE` | Apache License 2.0 and attribution notice |
+
+## Citation
+
+```bibtex
+@techreport{ottoboni2026lightpfn,
+  title  = {A Sling Against Giants: {LightPFN}, a 4.6M-parameter tabular in-context classifier designed to stay small},
+  author = {Ottoboni, Giorgio},
+  year   = {2026},
+  url    = {https://github.com/GioOtto/LightPFN}
+}
+```

artifacts.json

@@ -1,52 +0,0 @@
-{
-  "lightpfn-0.1.0-py3-none-any.whl": {
-    "sha256": "c97936566a9d82a012d865d40111263f4fea9f21eb8a44840ac4e040b5f55604",
-    "bytes": 46760,
-    "files": [
-      "lightpfn/__init__.py",
-      "lightpfn/checkpoint.py",
-      "lightpfn/device.py",
-      "lightpfn/pretrained.json",
-      "lightpfn/sklearn.py",
-      "lightpfn/model/__init__.py",
-      "lightpfn/model/layers.py",
-      "lightpfn/model/lightpfn.py",
-      "lightpfn/vulkan/__init__.py",
-      "lightpfn/vulkan/engine.py",
-      "lightpfn/vulkan/kernels.py",
-      "lightpfn/vulkan/model.py",
-      "lightpfn-0.1.0.dist-info/METADATA",
-      "lightpfn-0.1.0.dist-info/WHEEL",
-      "lightpfn-0.1.0.dist-info/licenses/LICENSE",
-      "lightpfn-0.1.0.dist-info/licenses/NOTICE",
-      "lightpfn-0.1.0.dist-info/RECORD"
-    ]
-  },
-  "lightpfn-0.1.0.tar.gz": {
-    "sha256": "3a9606af85ad6d178963717cb7ad0dec432ff44ea7bea000396bb34c0ff2ab5d",
-    "bytes": 45222,
-    "files": [
-      "lightpfn-0.1.0/docs/DEPENDENCIES.md",
-      "lightpfn-0.1.0/lightpfn/__init__.py",
-      "lightpfn-0.1.0/lightpfn/checkpoint.py",
-      "lightpfn-0.1.0/lightpfn/device.py",
-      "lightpfn-0.1.0/lightpfn/pretrained.json",
-      "lightpfn-0.1.0/lightpfn/sklearn.py",
-      "lightpfn-0.1.0/lightpfn/model/__init__.py",
-      "lightpfn-0.1.0/lightpfn/model/layers.py",
-      "lightpfn-0.1.0/lightpfn/model/lightpfn.py",
-      "lightpfn-0.1.0/lightpfn/vulkan/__init__.py",
-      "lightpfn-0.1.0/lightpfn/vulkan/engine.py",
-      "lightpfn-0.1.0/lightpfn/vulkan/kernels.py",
-      "lightpfn-0.1.0/lightpfn/vulkan/model.py",
-      "lightpfn-0.1.0/tests/test_release.py",
-      "lightpfn-0.1.0/tests/test_sklearn_api.py",
-      "lightpfn-0.1.0/.gitignore",
-      "lightpfn-0.1.0/LICENSE",
-      "lightpfn-0.1.0/NOTICE",
-      "lightpfn-0.1.0/pyproject.toml",
-      "lightpfn-0.1.0/docs/PACKAGE_README.md",
-      "lightpfn-0.1.0/PKG-INFO"
-    ]
-  }
-}

dependency_licenses.json

@@ -1,471 +0,0 @@
-{
-  "anyio": {
-    "version": "4.15.1",
-    "license_expression": "MIT",
-    "license_summary": "The MIT License (MIT)",
-    "license_classifiers": [],
-    "license_files": [
-      "anyio-4.15.1.dist-info/licenses/LICENSE"
-    ]
-  },
-  "cffi": {
-    "version": "2.1.1",
-    "license_expression": "MIT-0",
-    "license_summary": "Except when otherwise stated (look for LICENSE files in directories or",
-    "license_classifiers": [],
-    "license_files": [
-      "cffi-2.1.1.dist-info/licenses/LICENSE"
-    ]
-  },
-  "click": {
-    "version": "8.5.0",
-    "license_expression": "BSD-3-Clause",
-    "license_summary": "Copyright 2014 Pallets",
-    "license_classifiers": [],
-    "license_files": [
-      "click-8.5.0.dist-info/licenses/LICENSE.txt"
-    ]
-  },
-  "cloudpickle": {
-    "version": "3.1.2",
-    "license_expression": null,
-    "license_summary": "BSD-3-Clause",
-    "license_classifiers": [
-      "License :: OSI Approved :: BSD License"
-    ],
-    "license_files": [
-      "cloudpickle-3.1.2.dist-info/licenses/LICENSE"
-    ]
-  },
-  "colorama": {
-    "version": "0.4.6",
-    "license_expression": null,
-    "license_summary": "Copyright (c) 2010 Jonathan Hartley",
-    "license_classifiers": [
-      "License :: OSI Approved :: BSD License"
-    ],
-    "license_files": [
-      "colorama-0.4.6.dist-info/licenses/LICENSE.txt"
-    ]
-  },
-  "filelock": {
-    "version": "3.32.3",
-    "license_expression": "MIT",
-    "license_summary": "MIT License",
-    "license_classifiers": [
-      "License :: OSI Approved :: MIT License"
-    ],
-    "license_files": [
-      "filelock-3.32.3.dist-info/licenses/LICENSE"
-    ]
-  },
-  "fsspec": {
-    "version": "2026.7.0",
-    "license_expression": "BSD-3-Clause",
-    "license_summary": "BSD 3-Clause License",
-    "license_classifiers": [],
-    "license_files": [
-      "fsspec-2026.7.0.dist-info/licenses/LICENSE"
-    ]
-  },
-  "h11": {
-    "version": "0.16.0",
-    "license_expression": null,
-    "license_summary": "MIT",
-    "license_classifiers": [
-      "License :: OSI Approved :: MIT License"
-    ],
-    "license_files": [
-      "h11-0.16.0.dist-info/licenses/LICENSE.txt"
-    ]
-  },
-  "hf-xet": {
-    "version": "1.6.0",
-    "license_expression": "Apache-2.0",
-    "license_summary": "Apache License",
-    "license_classifiers": [
-      "License :: OSI Approved :: Apache Software License"
-    ],
-    "license_files": [
-      "hf_xet-1.6.0.dist-info/licenses/LICENSE"
-    ]
-  },
-  "httpcore2": {
-    "version": "2.13.1",
-    "license_expression": "BSD-3-Clause",
-    "license_summary": "Copyright \u00a9 2026 to present Pydantic Services Inc. and individual contributors.",
-    "license_classifiers": [
-      "License :: OSI Approved :: BSD License"
-    ],
-    "license_files": [
-      "httpcore2-2.13.1.dist-info/licenses/LICENSE.md"
-    ]
-  },
-  "httpx2": {
-    "version": "2.13.1",
-    "license_expression": "BSD-3-Clause",
-    "license_summary": "Copyright \u00a9 2026 to present Pydantic Services Inc. and individual contributors.",
-    "license_classifiers": [
-      "License :: OSI Approved :: BSD License"
-    ],
-    "license_files": [
-      "httpx2-2.13.1.dist-info/licenses/LICENSE.md"
-    ]
-  },
-  "huggingface-hub": {
-    "version": "2.1.1",
-    "license_expression": null,
-    "license_summary": "Apache-2.0",
-    "license_classifiers": [
-      "License :: OSI Approved :: Apache Software License"
-    ],
-    "license_files": [
-      "huggingface_hub-2.1.1.dist-info/licenses/LICENSE"
-    ]
-  },
-  "idna": {
-    "version": "3.20",
-    "license_expression": "BSD-3-Clause",
-    "license_summary": "BSD 3-Clause License",
-    "license_classifiers": [],
-    "license_files": [
-      "idna-3.20.dist-info/licenses/LICENSE.md"
-    ]
-  },
-  "jinja2": {
-    "version": "3.1.6",
-    "license_expression": null,
-    "license_summary": "Copyright 2007 Pallets",
-    "license_classifiers": [
-      "License :: OSI Approved :: BSD License"
-    ],
-    "license_files": [
-      "jinja2-3.1.6.dist-info/licenses/LICENSE.txt"
-    ]
-  },
-  "joblib": {
-    "version": "1.6.0",
-    "license_expression": "BSD-3-Clause",
-    "license_summary": "BSD 3-Clause License",
-    "license_classifiers": [],
-    "license_files": [
-      "joblib-1.6.0.dist-info/licenses/LICENSE.txt"
-    ]
-  },
-  "markupsafe": {
-    "version": "3.0.4",
-    "license_expression": "BSD-3-Clause",
-    "license_summary": "Copyright 2010 Pallets",
-    "license_classifiers": [],
-    "license_files": [
-      "markupsafe-3.0.4.dist-info/licenses/LICENSE.txt"
-    ]
-  },
-  "mpmath": {
-    "version": "1.3.0",
-    "license_expression": null,
-    "license_summary": "BSD",
-    "license_classifiers": [
-      "License :: OSI Approved :: BSD License"
-    ],
-    "license_files": [
-      "mpmath-1.3.0.dist-info/LICENSE"
-    ]
-  },
-  "narwhals": {
-    "version": "2.27.0",
-    "license_expression": "MIT",
-    "license_summary": "MIT License",
-    "license_classifiers": [],
-    "license_files": [
-      "narwhals-2.27.0.dist-info/licenses/LICENSE.md"
-    ]
-  },
-  "networkx": {
-    "version": "3.6.1",
-    "license_expression": "BSD-3-Clause",
-    "license_summary": "NetworkX is distributed with the 3-clause BSD license.",
-    "license_classifiers": [],
-    "license_files": [
-      "networkx-3.6.1.dist-info/licenses/LICENSE.txt"
-    ]
-  },
-  "numpy": {
-    "version": "2.5.3",
-    "license_expression": "BSD-3-Clause AND 0BSD AND MIT AND Zlib AND CC0-1.0",
-    "license_summary": "Copyright (c) 2005-2025, NumPy Developers.",
-    "license_classifiers": [],
-    "license_files": [
-      "numpy-2.5.3.dist-info/licenses/LICENSE.txt",
-      "numpy-2.5.3.dist-info/licenses/numpy/_core/include/numpy/libdivide/LICENSE.txt",
-      "numpy-2.5.3.dist-info/licenses/numpy/_core/src/common/pythoncapi-compat/COPYING",
-      "numpy-2.5.3.dist-info/licenses/numpy/_core/src/highway/LICENSE",
-      "numpy-2.5.3.dist-info/licenses/numpy/_core/src/multiarray/dragon4_LICENSE.txt",
-      "numpy-2.5.3.dist-info/licenses/numpy/_core/src/npysort/x86-simd-sort/LICENSE.md",
-      "numpy-2.5.3.dist-info/licenses/numpy/_core/src/umath/svml/LICENSE",
-      "numpy-2.5.3.dist-info/licenses/numpy/fft/pocketfft/LICENSE.md",
-      "numpy-2.5.3.dist-info/licenses/numpy/linalg/lapack_lite/LICENSE.txt",
-      "numpy-2.5.3.dist-info/licenses/numpy/ma/LICENSE",
-      "numpy-2.5.3.dist-info/licenses/numpy/random/LICENSE.md",
-      "numpy-2.5.3.dist-info/licenses/numpy/random/src/distributions/LICENSE.md",
-      "numpy-2.5.3.dist-info/licenses/numpy/random/src/mt19937/LICENSE.md",
-      "numpy-2.5.3.dist-info/licenses/numpy/random/src/pcg64/LICENSE.md",
-      "numpy-2.5.3.dist-info/licenses/numpy/random/src/philox/LICENSE.md",
-      "numpy-2.5.3.dist-info/licenses/numpy/random/src/sfc64/LICENSE.md",
-      "numpy-2.5.3.dist-info/licenses/numpy/random/src/splitmix64/LICENSE.md"
-    ]
-  },
-  "packaging": {
-    "version": "26.3",
-    "license_expression": "Apache-2.0 OR BSD-2-Clause",
-    "license_summary": "This software is made available under the terms of *either* of the licenses",
-    "license_classifiers": [],
-    "license_files": [
-      "packaging-26.3.dist-info/licenses/LICENSE",
-      "packaging-26.3.dist-info/licenses/LICENSE.APACHE",
-      "packaging-26.3.dist-info/licenses/LICENSE.BSD"
-    ]
-  },
-  "pycparser": {
-    "version": "3.0",
-    "license_expression": "BSD-3-Clause",
-    "license_summary": "pycparser -- A C parser in Python",
-    "license_classifiers": [],
-    "license_files": [
-      "pycparser-3.0.dist-info/licenses/LICENSE"
-    ]
-  },
-  "pyyaml": {
-    "version": "6.0.3",
-    "license_expression": null,
-    "license_summary": "MIT",
-    "license_classifiers": [
-      "License :: OSI Approved :: MIT License"
-    ],
-    "license_files": [
-      "pyyaml-6.0.3.dist-info/licenses/LICENSE"
-    ]
-  },
-  "rendercanvas": {
-    "version": "2.7.2",
-    "license_expression": null,
-    "license_summary": "BSD 2-Clause License",
-    "license_classifiers": [],
-    "license_files": [
-      "rendercanvas-2.7.2.dist-info/licenses/LICENSE"
-    ]
-  },
-  "safetensors": {
-    "version": "0.8.0",
-    "license_expression": null,
-    "license_summary": "Apache License",
-    "license_classifiers": [
-      "License :: OSI Approved :: Apache Software License"
-    ],
-    "license_files": [
-      "safetensors-0.8.0.dist-info/licenses/LICENSE"
-    ]
-  },
-  "scikit-learn": {
-    "version": "1.9.1",
-    "license_expression": "BSD-3-Clause",
-    "license_summary": "BSD 3-Clause License",
-    "license_classifiers": [],
-    "license_files": [
-      "scikit_learn-1.9.1.dist-info/licenses/COPYING"
-    ]
-  },
-  "scipy": {
-    "version": "1.18.1",
-    "license_expression": null,
-    "license_summary": "Copyright (c) 2001-2002 Enthought, Inc. 2003, SciPy Developers.",
-    "license_classifiers": [
-      "License :: OSI Approved :: BSD License"
-    ],
-    "license_files": [
-      "scipy-1.18.1.dist-info/LICENSE.txt"
-    ]
-  },
-  "setuptools": {
-    "version": "83.0.0",
-    "license_expression": "MIT",
-    "license_summary": "GNU LESSER GENERAL PUBLIC LICENSE",
-    "license_classifiers": [],
-    "license_files": [
-      "setuptools/_vendor/autocommand-2.2.2.dist-info/LICENSE",
-      "setuptools/_vendor/backports.tarfile-1.2.0.dist-info/LICENSE",
-      "setuptools/_vendor/importlib_metadata-8.7.1.dist-info/licenses/LICENSE",
-      "setuptools/_vendor/jaraco.text-4.0.0.dist-info/LICENSE",
-      "setuptools/_vendor/jaraco_context-6.1.0.dist-info/licenses/LICENSE",
-      "setuptools/_vendor/jaraco_functools-4.4.0.dist-info/licenses/LICENSE",
-      "setuptools/_vendor/more_itertools-10.8.0.dist-info/licenses/LICENSE",
-      "setuptools/_vendor/packaging-26.0.dist-info/licenses/LICENSE",
-      "setuptools/_vendor/packaging-26.0.dist-info/licenses/LICENSE.APACHE",
-      "setuptools/_vendor/packaging-26.0.dist-info/licenses/LICENSE.BSD",
-      "setuptools/_vendor/platformdirs-4.4.0.dist-info/licenses/LICENSE",
-      "setuptools/_vendor/tomli-2.4.0.dist-info/licenses/LICENSE",
-      "setuptools/_vendor/wheel-0.46.3.dist-info/licenses/LICENSE.txt",
-      "setuptools/_vendor/zipp-3.23.0.dist-info/licenses/LICENSE"
-    ]
-  },
-  "sympy": {
-    "version": "1.14.0",
-    "license_expression": null,
-    "license_summary": "BSD",
-    "license_classifiers": [
-      "License :: OSI Approved :: BSD License"
-    ],
-    "license_files": [
-      "sympy-1.14.0.dist-info/licenses/AUTHORS",
-      "sympy-1.14.0.dist-info/licenses/LICENSE"
-    ]
-  },
-  "threadpoolctl": {
-    "version": "3.7.0",
-    "license_expression": "BSD-3-Clause",
-    "license_summary": "Copyright (c) 2019, threadpoolctl contributors",
-    "license_classifiers": [],
-    "license_files": [
-      "threadpoolctl-3.7.0.dist-info/licenses/LICENSE"
-    ]
-  },
-  "torch": {
-    "version": "2.14.1+cpu",
-    "license_expression": "Apache-2.0 AND Apache-2.0 WITH LLVM-exception AND BSD-2-Clause AND BSD-3-Clause AND BSL-1.0 AND MIT",
-    "license_summary": "From PyTorch:",
-    "license_classifiers": [],
-    "license_files": [
-      "torch-2.14.1+cpu.dist-info/licenses/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/FP16/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/FXdiv/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/NNPACK/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/NVTX/LICENSE.txt",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/NVTX/docs/LICENSE.txt",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/NVTX/python/LICENSE.txt",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/NVTX/rust/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/NVTX/tools/docs/github-markdown-css/license",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/VulkanMemoryAllocator/LICENSE.txt",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/XNNPACK/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/aiter/3rdparty/composable_kernel/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/aiter/3rdparty/composable_kernel/docs/license.rst",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/aiter/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/benchmark/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/composable_kernel/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/composable_kernel/docs/license.rst",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/cpp-httplib/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/cpuinfo/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/cpuinfo/deps/clog/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/cudnn_frontend/LICENSE.txt",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/cutlass/LICENSE.txt",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/cutlass/python/LICENSE.txt",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/fbgemm/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/fbgemm/external/composable_kernel/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/fbgemm/external/composable_kernel/docs/license.rst",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/fbgemm/external/cpuinfo/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/fbgemm/external/cpuinfo/deps/clog/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/fbgemm/external/cutlass/LICENSE.txt",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/fbgemm/external/cutlass/python/LICENSE.txt",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/fbgemm/external/googletest/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/fbgemm/external/hipify_torch/LICENSE.txt",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/fbgemm/fbgemm_gpu/docs/src/general/License.rst",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/fbgemm/fbgemm_gpu/experimental/hstu/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/fbgemm/fbgemm_gpu/src/quantize_ops/mx/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/fbgemm/fbgemm_gpu/test/quantize/mx/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/flash-attention/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/flash-attention/csrc/composable_kernel/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/flash-attention/csrc/composable_kernel/docs/license.rst",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/flash-attention/csrc/cutlass/LICENSE.txt",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/flash-attention/csrc/cutlass/python/LICENSE.txt",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/flash-attention/flash_attn/cute/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/flash-attention/third_party/aiter/3rdparty/composable_kernel/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/flash-attention/third_party/aiter/3rdparty/composable_kernel/docs/license.rst",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/flash-attention/third_party/aiter/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/flatbuffers/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/flatbuffers/dart/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/flatbuffers/swift/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/fmt/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/gemmlowp/gemmlowp/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/gloo/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/googletest/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/ideep/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/ideep/mkl-dnn/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/ideep/mkl-dnn/third_party/gtest/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/ideep/mkl-dnn/third_party/opencl/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/kineto/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/kineto/libkineto/third_party/dynolog_headers/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/kineto/libkineto/third_party/fmt/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/kineto/libkineto/third_party/googletest/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/llvm-openmp/LICENSE.txt",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/mimalloc/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/miniz-3.0.2/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/mslk/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/mslk/external/composable_kernel/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/mslk/external/composable_kernel/docs/license.rst",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/mslk/external/cutlass/LICENSE.txt",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/mslk/external/cutlass/python/LICENSE.txt",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/mslk/external/googletest/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/mslk/external/hipify_torch/LICENSE.txt",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/mslk/mslk/attention/flash_attn/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/onnx/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/onnx/third_party/pybind11/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/perfetto/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/protobuf/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/protobuf/third_party/benchmark/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/protobuf/third_party/googletest/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/protobuf/third_party/googletest/googlemock/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/protobuf/third_party/googletest/googlemock/scripts/generator/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/protobuf/third_party/googletest/googletest/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/protobuf/third_party/utf8_range/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/psimd/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/pthreadpool/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/pybind11/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/python-peachpy/LICENSE.rst",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/sleef/LICENSE.txt",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/tensorpipe/LICENSE.txt",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/tensorpipe/third_party/googletest/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/tensorpipe/third_party/googletest/googlemock/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/tensorpipe/third_party/googletest/googlemock/scripts/generator/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/tensorpipe/third_party/googletest/googletest/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/tensorpipe/third_party/libnop/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/tensorpipe/third_party/libuv/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/tensorpipe/third_party/pybind11/LICENSE",
-      "torch-2.14.1+cpu.dist-info/licenses/third_party/tensorpipe/third_party/pybind11/tools/clang/LICENSE.TXT"
-    ]
-  },
-  "tqdm": {
-    "version": "4.70.1",
-    "license_expression": null,
-    "license_summary": "MPL-2.0 AND MIT",
-    "license_classifiers": [],
-    "license_files": [
-      "tqdm-4.70.1.dist-info/licenses/LICENCE"
-    ]
-  },
-  "truststore": {
-    "version": "0.10.4",
-    "license_expression": "MIT",
-    "license_summary": "The MIT License (MIT)",
-    "license_classifiers": [],
-    "license_files": [
-      "truststore-0.10.4.dist-info/licenses/LICENSE"
-    ]
-  },
-  "typing-extensions": {
-    "version": "4.16.0",
-    "license_expression": "PSF-2.0",
-    "license_summary": "A. HISTORY OF THE SOFTWARE",
-    "license_classifiers": [],
-    "license_files": [
-      "typing_extensions-4.16.0.dist-info/licenses/LICENSE"
-    ]
-  },
-  "wgpu": {
-    "version": "0.32.0",
-    "license_expression": null,
-    "license_summary": "BSD 2-Clause License",
-    "license_classifiers": [],
-    "license_files": [
-      "wgpu-0.32.0.dist-info/licenses/LICENSE"
-    ]
-  }
-}

dist/lightpfn-0.1.0.tar.gz

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:3a9606af85ad6d178963717cb7ad0dec432ff44ea7bea000396bb34c0ff2ab5d
-size 45222

installed_verification.json

@@ -1,17 +0,0 @@
-{
-  "package_version": "0.1.0",
-  "model": {
-    "repo_id": "ueuegio/LightPFN",
-    "revision": "bd389ab59a89dd0e05c9ecb7c642c08ee52e9637"
-  },
-  "installed_import": true,
-  "hub_download": true,
-  "offline_network_blocked": true,
-  "estimators": 4,
-  "max_probability_difference": 0.0,
-  "probability_shape": [
-    40,
-    3
-  ],
-  "finite": true
-}

package/.gitignore

@@ -1,53 +0,0 @@
-__pycache__/
-*.pyc
-dist/
-build/
-*.egg-info/
-.pytest_cache/
-runs/release/staging/
-runs/release/install_check/
-*.safetensors
-
-# datasets, synthetic pools, OpenML cache and TabICL offload files stay on the data disk
-data/
-
-# model checkpoints
-*.pt
-*.pt.tmp
-
-# scratch and local copies
-runs/snapshots/
-runs/codex_tmp/
-
-# full Codex transcripts (the prompts and final answers are tracked)
-runs/reviews/*.log
-runs/reviews/*.err
-
-# Codex patch dumps (the history is in git)
-runs/reviews/*.patch
-
-# bulky analysis intermediates
-runs/analysis_r1/*.npz
-runs/analysis_r1/prior_feature_stats.csv
-runs/analysis_r1/prior_geometry_all.csv
-runs/analysis_r1/matched_n128_feature_stats.csv
-
-# thermal watchdog output (the script is tracked)
-runs/thermal/temps.csv
-runs/thermal/watchdog.log
-runs/thermal/STOP
-
-# benchmark scratch (temporary runs, pools, test dirs)
-runs/bench_train/pytest_tmp/
-runs/bench_train/test_tmp/
-runs/bench_train/pools/
-runs/bench_train/runs/
-runs/bench_train/plan_*.json
-
-# LaTeX build files (the PDF is tracked)
-docs/report/*.aux
-docs/report/*.log
-docs/report/*.out
-
-# lock files of the evaluation queue
-runs/**/.eval.lock

package/LICENSE

@@ -1,202 +0,0 @@
-
-                                 Apache License
-                           Version 2.0, January 2004
-                        http://www.apache.org/licenses/
-
-   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
-
-   1. Definitions.
-
-      "License" shall mean the terms and conditions for use, reproduction,
-      and distribution as defined by Sections 1 through 9 of this document.
-
-      "Licensor" shall mean the copyright owner or entity authorized by
-      the copyright owner that is granting the License.
-
-      "Legal Entity" shall mean the union of the acting entity and all
-      other entities that control, are controlled by, or are under common
-      control with that entity. For the purposes of this definition,
-      "control" means (i) the power, direct or indirect, to cause the
-      direction or management of such entity, whether by contract or
-      otherwise, or (ii) ownership of fifty percent (50%) or more of the
-      outstanding shares, or (iii) beneficial ownership of such entity.
-
-      "You" (or "Your") shall mean an individual or Legal Entity
-      exercising permissions granted by this License.
-
-      "Source" form shall mean the preferred form for making modifications,
-      including but not limited to software source code, documentation
-      source, and configuration files.
-
-      "Object" form shall mean any form resulting from mechanical
-      transformation or translation of a Source form, including but
-      not limited to compiled object code, generated documentation,
-      and conversions to other media types.
-
-      "Work" shall mean the work of authorship, whether in Source or
-      Object form, made available under the License, as indicated by a
-      copyright notice that is included in or attached to the work
-      (an example is provided in the Appendix below).
-
-      "Derivative Works" shall mean any work, whether in Source or Object
-      form, that is based on (or derived from) the Work and for which the
-      editorial revisions, annotations, elaborations, or other modifications
-      represent, as a whole, an original work of authorship. For the purposes
-      of this License, Derivative Works shall not include works that remain
-      separable from, or merely link (or bind by name) to the interfaces of,
-      the Work and Derivative Works thereof.
-
-      "Contribution" shall mean any work of authorship, including
-      the original version of the Work and any modifications or additions
-      to that Work or Derivative Works thereof, that is intentionally
-      submitted to Licensor for inclusion in the Work by the copyright owner
-      or by an individual or Legal Entity authorized to submit on behalf of
-      the copyright owner. For the purposes of this definition, "submitted"
-      means any form of electronic, verbal, or written communication sent
-      to the Licensor or its representatives, including but not limited to
-      communication on electronic mailing lists, source code control systems,
-      and issue tracking systems that are managed by, or on behalf of, the
-      Licensor for the purpose of discussing and improving the Work, but
-      excluding communication that is conspicuously marked or otherwise
-      designated in writing by the copyright owner as "Not a Contribution."
-
-      "Contributor" shall mean Licensor and any individual or Legal Entity
-      on behalf of whom a Contribution has been received by Licensor and
-      subsequently incorporated within the Work.
-
-   2. Grant of Copyright License. Subject to the terms and conditions of
-      this License, each Contributor hereby grants to You a perpetual,
-      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
-      copyright license to reproduce, prepare Derivative Works of,
-      publicly display, publicly perform, sublicense, and distribute the
-      Work and such Derivative Works in Source or Object form.
-
-   3. Grant of Patent License. Subject to the terms and conditions of
-      this License, each Contributor hereby grants to You a perpetual,
-      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
-      (except as stated in this section) patent license to make, have made,
-      use, offer to sell, sell, import, and otherwise transfer the Work,
-      where such license applies only to those patent claims licensable
-      by such Contributor that are necessarily infringed by their
-      Contribution(s) alone or by combination of their Contribution(s)
-      with the Work to which such Contribution(s) was submitted. If You
-      institute patent litigation against any entity (including a
-      cross-claim or counterclaim in a lawsuit) alleging that the Work
-      or a Contribution incorporated within the Work constitutes direct
-      or contributory patent infringement, then any patent licenses
-      granted to You under this License for that Work shall terminate
-      as of the date such litigation is filed.
-
-   4. Redistribution. You may reproduce and distribute copies of the
-      Work or Derivative Works thereof in any medium, with or without
-      modifications, and in Source or Object form, provided that You
-      meet the following conditions:
-
-      (a) You must give any other recipients of the Work or
-          Derivative Works a copy of this License; and
-
-      (b) You must cause any modified files to carry prominent notices
-          stating that You changed the files; and
-
-      (c) You must retain, in the Source form of any Derivative Works
-          that You distribute, all copyright, patent, trademark, and
-          attribution notices from the Source form of the Work,
-          excluding those notices that do not pertain to any part of
-          the Derivative Works; and
-
-      (d) If the Work includes a "NOTICE" text file as part of its
-          distribution, then any Derivative Works that You distribute must
-          include a readable copy of the attribution notices contained
-          within such NOTICE file, excluding those notices that do not
-          pertain to any part of the Derivative Works, in at least one
-          of the following places: within a NOTICE text file distributed
-          as part of the Derivative Works; within the Source form or
-          documentation, if provided along with the Derivative Works; or,
-          within a display generated by the Derivative Works, if and
-          wherever such third-party notices normally appear. The contents
-          of the NOTICE file are for informational purposes only and
-          do not modify the License. You may add Your own attribution
-          notices within Derivative Works that You distribute, alongside
-          or as an addendum to the NOTICE text from the Work, provided
-          that such additional attribution notices cannot be construed
-          as modifying the License.
-
-      You may add Your own copyright statement to Your modifications and
-      may provide additional or different license terms and conditions
-      for use, reproduction, or distribution of Your modifications, or
-      for any such Derivative Works as a whole, provided Your use,
-      reproduction, and distribution of the Work otherwise complies with
-      the conditions stated in this License.
-
-   5. Submission of Contributions. Unless You explicitly state otherwise,
-      any Contribution intentionally submitted for inclusion in the Work
-      by You to the Licensor shall be under the terms and conditions of
-      this License, without any additional terms or conditions.
-      Notwithstanding the above, nothing herein shall supersede or modify
-      the terms of any separate license agreement you may have executed
-      with Licensor regarding such Contributions.
-
-   6. Trademarks. This License does not grant permission to use the trade
-      names, trademarks, service marks, or product names of the Licensor,
-      except as required for reasonable and customary use in describing the
-      origin of the Work and reproducing the content of the NOTICE file.
-
-   7. Disclaimer of Warranty. Unless required by applicable law or
-      agreed to in writing, Licensor provides the Work (and each
-      Contributor provides its Contributions) on an "AS IS" BASIS,
-      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
-      implied, including, without limitation, any warranties or conditions
-      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
-      PARTICULAR PURPOSE. You are solely responsible for determining the
-      appropriateness of using or redistributing the Work and assume any
-      risks associated with Your exercise of permissions under this License.
-
-   8. Limitation of Liability. In no event and under no legal theory,
-      whether in tort (including negligence), contract, or otherwise,
-      unless required by applicable law (such as deliberate and grossly
-      negligent acts) or agreed to in writing, shall any Contributor be
-      liable to You for damages, including any direct, indirect, special,
-      incidental, or consequential damages of any character arising as a
-      result of this License or out of the use or inability to use the
-      Work (including but not limited to damages for loss of goodwill,
-      work stoppage, computer failure or malfunction, or any and all
-      other commercial damages or losses), even if such Contributor
-      has been advised of the possibility of such damages.
-
-   9. Accepting Warranty or Additional Liability. While redistributing
-      the Work or Derivative Works thereof, You may choose to offer,
-      and charge a fee for, acceptance of support, warranty, indemnity,
-      or other liability obligations and/or rights consistent with this
-      License. However, in accepting such obligations, You may act only
-      on Your own behalf and on Your sole responsibility, not on behalf
-      of any other Contributor, and only if You agree to indemnify,
-      defend, and hold each Contributor harmless for any liability
-      incurred by, or claims asserted against, such Contributor by reason
-      of your accepting any such warranty or additional liability.
-
-   END OF TERMS AND CONDITIONS
-
-   APPENDIX: How to apply the Apache License to your work.
-
-      To apply the Apache License to your work, attach the following
-      boilerplate notice, with the fields enclosed by brackets "[]"
-      replaced with your own identifying information. (Don't include
-      the brackets!)  The text should be enclosed in the appropriate
-      comment syntax for the file format. We also recommend that a
-      file or class name and description of purpose be included on the
-      same "printed page" as the copyright notice for easier
-      identification within third-party archives.
-
-   Copyright [yyyy] [name of copyright owner]
-
-   Licensed under the Apache License, Version 2.0 (the "License");
-   you may not use this file except in compliance with the License.
-   You may obtain a copy of the License at
-
-       http://www.apache.org/licenses/LICENSE-2.0
-
-   Unless required by applicable law or agreed to in writing, software
-   distributed under the License is distributed on an "AS IS" BASIS,
-   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-   See the License for the specific language governing permissions and
-   limitations under the License.

package/NOTICE

@@ -1,6 +0,0 @@
-LightPFN
-Copyright 2026 Giorgio (GioOtto)
-
-Code and released model weights are licensed under the Apache License, Version 2.0.
-The model was trained from scratch using synthetic data, without distillation
-or training on weights or outputs from other tabular foundation models.

package/PKG-INFO

@@ -1,123 +0,0 @@
-Metadata-Version: 2.5
-Name: lightpfn
-Version: 0.1.0
-Summary: A compact tabular classifier pretrained only on synthetic data
-Project-URL: Repository, https://github.com/GioOtto/gioPFN
-Project-URL: Models, https://huggingface.co/ueuegio/LightPFN
-Author-email: Giorgio <247403232+GioOtto@users.noreply.github.com>
-License-Expression: Apache-2.0
-License-File: LICENSE
-License-File: NOTICE
-Classifier: Development Status :: 3 - Alpha
-Classifier: Intended Audience :: Science/Research
-Classifier: Operating System :: OS Independent
-Classifier: Programming Language :: Python :: 3
-Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
-Requires-Python: >=3.10
-Requires-Dist: numpy>=1.24
-Requires-Dist: torch>=2.6
-Provides-Extra: dev
-Requires-Dist: build>=1.2; extra == 'dev'
-Requires-Dist: hatchling>=1.27; extra == 'dev'
-Requires-Dist: pytest>=8; extra == 'dev'
-Provides-Extra: hf
-Requires-Dist: huggingface-hub>=0.27; extra == 'hf'
-Requires-Dist: safetensors>=0.5; extra == 'hf'
-Provides-Extra: sklearn
-Requires-Dist: scikit-learn>=1.6; extra == 'sklearn'
-Provides-Extra: vulkan
-Requires-Dist: wgpu<0.33,>=0.32; extra == 'vulkan'
-Description-Content-Type: text/markdown
-
-# LightPFN
-
-A **4,603,088-parameter** in-context classifier pretrained from scratch only on
-synthetic tabular tasks. `fit()` builds an inference context; no gradient training
-is performed on your dataset. Code and released weights: **Apache 2.0**.
-
-Python 3.10+; CPU, PyTorch CUDA/ROCm, and optional Vulkan inference.
-The base package requires only **torch and numpy**. Synthetic priors, training
-code, evaluation harnesses, and third-party foundation models are excluded from
-both distribution artifacts.
-
-## Install
-
-Install a supplied wheel, adding the sklearn and Hugging Face extras:
-
-```bash
-pip install "./lightpfn-0.1.0-py3-none-any.whl[sklearn,hf]"
-hf auth login
-```
-
-The initial Hugging Face repository is private: an authorized account/token is
-required for downloads. Authentication is handled by `huggingface_hub`, using
-its cached login or `HF_TOKEN`; credentials are never stored in the package.
-The package is not yet published on PyPI.
-
-From the source tree: `pip install ".[sklearn,hf]"`. For Vulkan add the `vulkan`
-extra. For a CPU installation, install a CPU PyTorch wheel from the official
-PyTorch index first; installing torch from PyPI may pull GPU runtime packages.
-
-## Use
-
-```python
-from lightpfn import LightPFNClassifier
-
-clf = LightPFNClassifier(device="cpu", n_estimators=4, random_state=0)
-clf.fit(X_train, y_train)
-probabilities = clf.predict_proba(X_test)  # columns follow clf.classes_
-predictions = clf.predict(X_test)
-```
-
-The constructor is cheap. The first `fit()` loads the released model from an
-immutable Hugging Face commit, and subsequent use benefits from the HF cache.
-`local_files_only=True` prohibits network downloads. `checkpoint=` accepts a
-local safetensors directory/file or a compatible legacy `.pt` checkpoint.
-`repo_id=` and `revision=` override the released model; custom repositories
-require an explicit revision. No remote Python code is loaded.
-
-The estimator inherits `ClassifierMixin` and `BaseEstimator`, supports
-`get_params`, `set_params`, cloning, `Pipeline`, `GridSearchCV`, `score`, feature
-name validation and `n_features_in_`. Use `random_state` or the legacy `seed`.
-Standard estimator checks are exercised with local weights. Two strict query
-invariance checks have documented FP32 tolerances: changing row order or batch
-size can change probabilities around 1e-7. For cloning/model selection prefer
-`checkpoint=` or the default pretrained model; sklearn's parameter hash check
-does not reliably compare raw torch modules because it hashes storage identity.
-Input must be a dense numeric table; NaN values are supported, infinity and
-sparse tables are rejected. Encode strings/categories before fitting, for
-example using sklearn's `OrdinalEncoder` with missing/unknown values mapped to
-NaN. Native categorical semantics are not implemented. The deprecated `cat`
-mask does not change predictions and warns when categorical columns are marked.
-
-The model was trained for **2-10 classes**. A single-class fit returns its constant
-class probability. For more than 10 classes, use a different classifier.
-Above `max_context` (default 20,000), each member uses a stratified subset that
-keeps every class. Larger contexts increase compute and memory substantially.
-`n_threads` sets PyTorch's process-wide CPU thread count. On Vulkan, `auto`
-falls back to CPU if the required buffers exceed adapter limits.
-
-Low-level use with only torch/numpy:
-
-```python
-from lightpfn import load_model
-
-model = load_model("inference.pt")  # restricted weights_only=True
-```
-
-To load/export safetensors, install the `hf` extra and use `load_model(folder)`
-or `save_model(model, folder)`. The release contains **safetensors + JSON**;
-unrestricted pickle loading is never used by the inference package.
-
-## Evidence and limitations
-
-The final `final_long` checkpoint follows 120,000 base steps and 39,250
-long-context steps. Selection used held-out synthetic tasks, probes and 55
-OpenML-CC18 tasks outside TabArena. TabArena was inspected only for finalists.
-On 38 classification tasks using official splits, first repeat, four estimators
-achieved mean AUC 0.8581 and lower task error than default CatBoost on 76% of
-tasks. This is a local evaluation, not an official TabArena submission, and does
-not establish superiority over tuned GBDTs. Categorical-heavy datasets remain
-a weakness; regression is unsupported.
-
-The full experiment history and technical report are in the source repository.

package/docs/DEPENDENCIES.md

@@ -1,28 +0,0 @@
-# Inference dependency licenses
-
-LightPFN's code and released weights use Apache-2.0. Dependency packages are
-installed separately and retain their upstream licenses; their bundled notices
-and licenses must remain intact if you redistribute those packages.
-
-| Dependency | Role | Upstream license |
-|---|---|---|
-| numpy | Required, numeric arrays | BSD-3-Clause; wheels also include permissive third-party notices |
-| torch | Required, model execution and restricted local loader | BSD-3-Clause; tested CPU wheel includes Apache-2.0, LLVM exception, BSD, BSL-1.0, MIT notices |
-| scikit-learn | Optional `sklearn` extra | BSD-3-Clause |
-| huggingface_hub | Optional `hf` extra | Apache-2.0 |
-| safetensors | Optional `hf` extra | Apache-2.0 |
-| wgpu | Optional `vulkan` extra | BSD-2-Clause; native wgpu components have their own permissive notices |
-
-These upstream licenses permit use with this Apache-2.0 package. They do not
-relicense the dependencies as Apache-2.0. No other tabular foundation model is
-a dependency of either the wheel or source distribution. Synthetic generation,
-training and benchmark modules remain in the research repository and are not
-distributed by this inference package.
-
-`runs/release/dependency_licenses.json` records the resolved inference dependency
-closure, versions and license metadata from the environment used to validate
-the release. Future dependency resolutions should be audited separately.
-The 35-package validation environment uses permissive BSD/MIT/Apache/PSF
-licenses, together with tqdm's MPL-2.0/MIT license. Its rendercanvas license was
-verified from the installed BSD-2-Clause license file; scipy retains its BSD
-license and the third-party notices bundled in its wheel.

package/docs/PACKAGE_README.md

@@ -1,92 +0,0 @@
-# LightPFN
-
-A **4,603,088-parameter** in-context classifier pretrained from scratch only on
-synthetic tabular tasks. `fit()` builds an inference context; no gradient training
-is performed on your dataset. Code and released weights: **Apache 2.0**.
-
-Python 3.10+; CPU, PyTorch CUDA/ROCm, and optional Vulkan inference.
-The base package requires only **torch and numpy**. Synthetic priors, training
-code, evaluation harnesses, and third-party foundation models are excluded from
-both distribution artifacts.
-
-## Install
-
-Install a supplied wheel, adding the sklearn and Hugging Face extras:
-
-```bash
-pip install "./lightpfn-0.1.0-py3-none-any.whl[sklearn,hf]"
-hf auth login
-```
-
-The initial Hugging Face repository is private: an authorized account/token is
-required for downloads. Authentication is handled by `huggingface_hub`, using
-its cached login or `HF_TOKEN`; credentials are never stored in the package.
-The package is not yet published on PyPI.
-
-From the source tree: `pip install ".[sklearn,hf]"`. For Vulkan add the `vulkan`
-extra. For a CPU installation, install a CPU PyTorch wheel from the official
-PyTorch index first; installing torch from PyPI may pull GPU runtime packages.
-
-## Use
-
-```python
-from lightpfn import LightPFNClassifier
-
-clf = LightPFNClassifier(device="cpu", n_estimators=4, random_state=0)
-clf.fit(X_train, y_train)
-probabilities = clf.predict_proba(X_test)  # columns follow clf.classes_
-predictions = clf.predict(X_test)
-```
-
-The constructor is cheap. The first `fit()` loads the released model from an
-immutable Hugging Face commit, and subsequent use benefits from the HF cache.
-`local_files_only=True` prohibits network downloads. `checkpoint=` accepts a
-local safetensors directory/file or a compatible legacy `.pt` checkpoint.
-`repo_id=` and `revision=` override the released model; custom repositories
-require an explicit revision. No remote Python code is loaded.
-
-The estimator inherits `ClassifierMixin` and `BaseEstimator`, supports
-`get_params`, `set_params`, cloning, `Pipeline`, `GridSearchCV`, `score`, feature
-name validation and `n_features_in_`. Use `random_state` or the legacy `seed`.
-Standard estimator checks are exercised with local weights. Two strict query
-invariance checks have documented FP32 tolerances: changing row order or batch
-size can change probabilities around 1e-7. For cloning/model selection prefer
-`checkpoint=` or the default pretrained model; sklearn's parameter hash check
-does not reliably compare raw torch modules because it hashes storage identity.
-Input must be a dense numeric table; NaN values are supported, infinity and
-sparse tables are rejected. Encode strings/categories before fitting, for
-example using sklearn's `OrdinalEncoder` with missing/unknown values mapped to
-NaN. Native categorical semantics are not implemented. The deprecated `cat`
-mask does not change predictions and warns when categorical columns are marked.
-
-The model was trained for **2-10 classes**. A single-class fit returns its constant
-class probability. For more than 10 classes, use a different classifier.
-Above `max_context` (default 20,000), each member uses a stratified subset that
-keeps every class. Larger contexts increase compute and memory substantially.
-`n_threads` sets PyTorch's process-wide CPU thread count. On Vulkan, `auto`
-falls back to CPU if the required buffers exceed adapter limits.
-
-Low-level use with only torch/numpy:
-
-```python
-from lightpfn import load_model
-
-model = load_model("inference.pt")  # restricted weights_only=True
-```
-
-To load/export safetensors, install the `hf` extra and use `load_model(folder)`
-or `save_model(model, folder)`. The release contains **safetensors + JSON**;
-unrestricted pickle loading is never used by the inference package.
-
-## Evidence and limitations
-
-The final `final_long` checkpoint follows 120,000 base steps and 39,250
-long-context steps. Selection used held-out synthetic tasks, probes and 55
-OpenML-CC18 tasks outside TabArena. TabArena was inspected only for finalists.
-On 38 classification tasks using official splits, first repeat, four estimators
-achieved mean AUC 0.8581 and lower task error than default CatBoost on 76% of
-tasks. This is a local evaluation, not an official TabArena submission, and does
-not establish superiority over tuned GBDTs. Categorical-heavy datasets remain
-a weakness; regression is unsupported.
-
-The full experiment history and technical report are in the source repository.

package/lightpfn/__init__.py

@@ -1,25 +0,0 @@
-"""LightPFN: compact tabular inference, pretrained only on synthetic data."""
-
-__version__ = "0.1.0"
-__all__ = ["Config", "LightPFN", "LightPFNClassifier", "load_model", "load_pretrained", "save_model", "__version__"]
-
-
-def __getattr__(name):
-    if name in ("Config", "LightPFN"):
-        from lightpfn.model import lightpfn
-        value = getattr(lightpfn, name)
-    elif name == "LightPFNClassifier":
-        try:
-            from lightpfn.sklearn import LightPFNClassifier
-        except ModuleNotFoundError as exc:
-            if exc.name != "sklearn":
-                raise
-            raise ImportError("LightPFNClassifier requires `pip install lightpfn[sklearn]`.") from exc
-        value = LightPFNClassifier
-    elif name in ("load_model", "load_pretrained", "save_model"):
-        from lightpfn import checkpoint
-        value = getattr(checkpoint, name)
-    else:
-        raise AttributeError(name)
-    globals()[name] = value
-    return value

package/lightpfn/checkpoint.py

@@ -1,107 +0,0 @@
-"""Safe local checkpoints and revision-pinned Hugging Face downloads.
-
-The release format is model.safetensors plus config.json. Legacy tensor/dict
-checkpoints remain supported using PyTorch's restricted weights-only loader.
-There is deliberately no fallback to unrestricted pickle.
-"""
-
-import json
-from collections.abc import Mapping
-from dataclasses import asdict
-from importlib.resources import files
-from pathlib import Path
-
-import torch
-
-from lightpfn.model.lightpfn import Config, LightPFN
-
-
-def _safetensors():
-    try:
-        from safetensors.torch import load_file, save_file
-    except ImportError as exc:
-        raise ImportError("Safetensors weights require `pip install lightpfn[hf]`.") from exc
-    return load_file, save_file
-
-
-def _config(values):
-    if not isinstance(values, Mapping):
-        raise ValueError("Checkpoint config must be a mapping of architecture parameters.")
-    values = dict(values)
-    if "group_offsets" in values:
-        values["group_offsets"] = tuple(values["group_offsets"])
-    return Config(**values)
-
-
-def load_model(path, device="cpu"):
-    """Load a safetensors directory/file or a weights-only compatible .pt file.
-
-    Safetensors files need config.json in the same directory. Full training
-    checkpoints containing arbitrary Python objects are intentionally rejected.
-    """
-    path = Path(path)
-    if path.is_dir():
-        path = path / "model.safetensors"
-    if path.suffix == ".safetensors":
-        load_file, _ = _safetensors()
-        config = json.loads(path.with_name("config.json").read_text(encoding="utf-8"))
-        state = load_file(str(path), device="cpu")
-    else:
-        checkpoint = torch.load(path, map_location="cpu", weights_only=True)
-        if not isinstance(checkpoint, Mapping) or "config" not in checkpoint:
-            raise ValueError("Checkpoint must contain 'config' and 'ema' or 'model' weights.")
-        config = checkpoint["config"]
-        state = checkpoint.get("ema", checkpoint.get("model"))
-    if not isinstance(state, Mapping) or not state or not all(
-        isinstance(key, str) and isinstance(value, torch.Tensor) for key, value in state.items()
-    ):
-        raise ValueError("Checkpoint weights must be a nonempty tensor state dictionary.")
-    model = LightPFN(_config(config))
-    model.load_state_dict(state, strict=True)
-    return model.to(device).eval()
-
-
-def save_model(model, directory):
-    """Export unfused inference weights and JSON config, excluding optimizer state."""
-    if getattr(model, "is_folded", False):
-        raise ValueError("Export the original model; folded weights use a different layout.")
-    _, save_file = _safetensors()
-    directory = Path(directory)
-    directory.mkdir(parents=True, exist_ok=True)
-    state = {key: value.detach().cpu().contiguous().clone() for key, value in model.state_dict().items()}
-    save_file(state, str(directory / "model.safetensors"), metadata={"format": "pt"})
-    (directory / "config.json").write_text(json.dumps(asdict(model.cfg), indent=2) + "\n", encoding="utf-8")
-    return directory
-
-
-def pretrained_spec():
-    """Return the packaged model repository and immutable revision."""
-    return json.loads(files("lightpfn").joinpath("pretrained.json").read_text(encoding="utf-8"))
-
-
-def load_pretrained(*, repo_id=None, revision=None, device="cpu", cache_dir=None, local_files_only=False):
-    """Load the released model, or an explicit repo and revision using cached HF authentication.
-
-    A custom repository requires a revision. No remote Python code is executed.
-    For private repositories authenticate with `hf auth login` or HF_TOKEN.
-    """
-    try:
-        from huggingface_hub import hf_hub_download
-    except ImportError as exc:
-        raise ImportError("Hugging Face downloads require `pip install lightpfn[hf]`.") from exc
-    spec = pretrained_spec()
-    repo_id = spec["repo_id"] if repo_id is None else repo_id
-    if revision is None:
-        if repo_id != spec["repo_id"]:
-            raise ValueError("Provide revision when using a custom Hugging Face repository.")
-        revision = spec["revision"]
-    if not revision:
-        raise ValueError("No pretrained revision configured; provide a checkpoint or explicit revision.")
-    options = dict(repo_id=repo_id, revision=revision, cache_dir=cache_dir, local_files_only=local_files_only)
-    config_path = Path(hf_hub_download(filename="config.json", **options))
-    # Use the resolved commit from the snapshot path for both files, even when a
-    # caller explicitly chose a mutable branch or tag.
-    resolved = config_path.parent.name
-    options["revision"] = resolved
-    weights_path = Path(hf_hub_download(filename="model.safetensors", **options))
-    return load_model(weights_path, device=device)

package/lightpfn/device.py

@@ -1,51 +0,0 @@
-"""Inference device: "auto" takes torch's GPU backend when there is one (CUDA on NVIDIA, also ROCm builds of
-torch, which use the same "cuda" device), otherwise a GPU with a Vulkan driver (AMD, Intel or NVIDIA,
-through wgpu), otherwise the CPU. Each can be asked for explicitly: "cuda[:i]", "vulkan[:i]" (index into
-lightpfn.vulkan.adapters()), "cpu". The environment variable LIGHTPFN_DEVICE replaces "auto" with any of
-these, e.g. LIGHTPFN_DEVICE=cpu to keep a GPU free.
-"""
-
-import os
-
-import torch
-
-DEVICE_ENV = "LIGHTPFN_DEVICE"
-
-
-def resolve_device(device="auto"):
-    """The device string LightPFNClassifier runs on: "cuda[:i]", "vulkan[:i]", "cpu" (or "mps")."""
-    device = "auto" if device is None else str(device).strip().lower()
-    if device == "auto":
-        device = os.environ.get(DEVICE_ENV, "").strip().lower() or "auto"
-    if device == "auto":
-        if torch.cuda.is_available():
-            return "cuda"
-        from lightpfn import vulkan
-
-        return "vulkan" if vulkan.is_available() else "cpu"
-    kind, _, index = device.partition(":")
-    if ":" in device and not index:
-        raise ValueError(f"device {device!r}: missing index after ':'")
-    if index and not index.isdigit():
-        raise ValueError(f"device {device!r}: the index after ':' must be a number")
-    if kind == "cpu":
-        return "cpu"
-    if kind == "cuda":
-        if not torch.cuda.is_available():
-            raise RuntimeError("device 'cuda' requested, but this torch build sees no CUDA/ROCm GPU")
-        return device
-    if kind == "vulkan":
-        from lightpfn import vulkan
-
-        if not vulkan.is_available(int(index) if index else None):
-            raise RuntimeError(f"device {device!r} requested, but no Vulkan adapter was found: it needs a GPU "
-                               "driver with Vulkan and `pip install wgpu` (adapters: lightpfn.vulkan.adapters())")
-        return device
-    if kind == "mps":
-        return device
-    raise ValueError(f"unknown device {device!r}: use 'auto', 'cpu', 'cuda[:i]' or 'vulkan[:i]'")
-
-
-def kind(device):
-    """"cpu", "cuda", "vulkan" or "mps" of a resolved device string."""
-    return str(device).partition(":")[0]

package/lightpfn/model/__init__.py

@@ -1 +0,0 @@
-"""LightPFN architecture; the inference package does not import synthetic priors."""

package/lightpfn/model/layers.py

@@ -1,190 +0,0 @@
-"""Building blocks shared by the three stages of the model.
-
-Tensors are (batch, sequence, dim) unless a name says otherwise. Attention runs through
-F.scaled_dot_product_attention, so the same code uses flash kernels on GPU and fused CPU kernels.
-"""
-
-import math
-
-import torch
-import torch.nn.functional as F
-from torch import nn
-
-
-class MLP(nn.Module):
-    """Two-layer GELU feed-forward block; the output projection starts at zero so every residual
-    block starts as the identity (as in TabPFN-3)."""
-
-    def __init__(self, dim, hidden):
-        super().__init__()
-        self.fc1 = nn.Linear(dim, hidden, bias=False)
-        self.fc2 = nn.Linear(hidden, dim, bias=False)
-        nn.init.zeros_(self.fc2.weight)
-
-    def forward(self, x):
-        return self.fc2(F.gelu(self.fc1(x)))
-
-
-class SoftmaxScaling(nn.Module):
-    """Scales attention queries as a function of the number of keys n (TabPFN-3 SoftmaxScalingMLP):
-    q * base(log n) * (1 + tanh(mod(q))). Keeps attention sharp when the context grows far beyond
-    the lengths seen in training. base starts at 1 so the layer starts as a no-op."""
-
-    def __init__(self, n_heads, head_dim, hidden=64):
-        super().__init__()
-        self.n_heads, self.head_dim = n_heads, head_dim
-        self.base = nn.Sequential(nn.Linear(1, hidden), nn.GELU(), nn.Linear(hidden, n_heads * head_dim))
-        self.mod = nn.Sequential(nn.Linear(head_dim, hidden), nn.GELU(), nn.Linear(hidden, head_dim))
-        nn.init.zeros_(self.base[2].weight)
-        nn.init.ones_(self.base[2].bias)
-        nn.init.zeros_(self.mod[2].weight)
-        nn.init.zeros_(self.mod[2].bias)
-
-    def forward(self, q, n):
-        """q: (B, H, L, D) queries, n: number of keys."""
-        logn = torch.full((1, 1), math.log(max(n, 2)), device=q.device, dtype=q.dtype)
-        base = self.base(logn).view(1, self.n_heads, 1, self.head_dim)
-        return q * base * (1 + torch.tanh(self.mod(q)))
-
-
-# CUDA attention kernels put the batch on a grid dimension of at most 65535 blocks: a larger batch fails
-# with "invalid argument" (the row stages of a 2-task group of 50k-row tables have 100k rows in the batch).
-SDPA_MAX_BATCH = 32768
-
-
-def sdpa(q, k, v, mask=None):
-    """F.scaled_dot_product_attention in chunks of at most SDPA_MAX_BATCH along the batch dim (a size-1
-    batch dim of k, v or mask broadcasts)."""
-    B = q.shape[0]
-    if B <= SDPA_MAX_BATCH:
-        return F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
-    out = []
-    for i in range(0, B, SDPA_MAX_BATCH):
-        part = [t if t is None or t.shape[0] == 1 else t[i:i + SDPA_MAX_BATCH] for t in (k, v, mask)]
-        out.append(F.scaled_dot_product_attention(q[i:i + SDPA_MAX_BATCH], *part[:2], attn_mask=part[2]))
-    return torch.cat(out)
-
-
-def rope(x, pos, base=10000.0, interleaved=False):
-    """Rotary position embedding (rotate-half form) on the last dim of x: (..., L, D) with
-    positions pos: (L,). interleaved=True rotates the pairs (2i, 2i + 1) instead of (i, i + D/2),
-    with one complex multiply: the same function for weights passed through interleave_rotary()."""
-    d = x.shape[-1]
-    inv_freq = base ** (-torch.arange(0, d, 2, device=x.device, dtype=torch.float32) / d)
-    ang = pos.to(torch.float32)[:, None] * inv_freq[None]  # (L, D/2)
-    if interleaved and x.dtype in (torch.float32, torch.float64):
-        rot = torch.complex(ang.cos(), ang.sin()).to(torch.complex64 if x.dtype == torch.float32 else torch.complex128)
-        return torch.view_as_real(torch.view_as_complex(x.unflatten(-1, (d // 2, 2))) * rot).flatten(-2)
-    cos, sin = ang.cos().to(x.dtype), ang.sin().to(x.dtype)
-    if interleaved:  # reduced precision (autocast): the same arithmetic as below on adjacent pairs
-        x1, x2 = x[..., 0::2], x[..., 1::2]
-        return torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1).flatten(-2)
-    x1, x2 = x[..., : d // 2], x[..., d // 2 :]
-    return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
-
-
-class Attention(nn.Module):
-    """Multi-head attention with separate query and key/value inputs.
-
-    kv_heads_query < n_heads lets a set of queries use only the first key/value heads (multi-query
-    attention for test rows, as in TabPFN-3.5): keys/values of the other heads never need to be
-    cached for prediction."""
-
-    def __init__(self, dim, n_heads, scaling=False):
-        super().__init__()
-        assert dim % n_heads == 0
-        self.n_heads, self.head_dim = n_heads, dim // n_heads
-        self.q = nn.Linear(dim, dim, bias=False)
-        self.kv = nn.Linear(dim, 2 * dim, bias=False)
-        self.out = nn.Linear(dim, dim, bias=False)
-        nn.init.zeros_(self.out.weight)
-        self.scaling = SoftmaxScaling(n_heads, self.head_dim) if scaling else None
-
-    def split(self, x):
-        B, L, _ = x.shape
-        return x.view(B, L, self.n_heads, self.head_dim).transpose(1, 2)  # (B, H, L, D)
-
-    def keys_values(self, x):
-        k, v = self.kv(x).chunk(2, dim=-1)
-        return self.split(k), self.split(v)
-
-    def attend(self, x, k, v, mask=None, q_rope=None, kv_heads=None):
-        """x: (B, Lq, dim) queries; k, v: (B, Hkv, Lk, D). kv_heads=h: all query heads use the
-        first h key/value heads."""
-        q = self.split(self.q(x))
-        if q_rope is not None:
-            q = q_rope(q)
-        if self.scaling is not None:
-            q = self.scaling(q, k.shape[2])
-        if kv_heads is not None:
-            k = k[:, :kv_heads].repeat_interleave(self.n_heads // kv_heads, dim=1)
-            v = v[:, :kv_heads].repeat_interleave(self.n_heads // kv_heads, dim=1)
-        o = sdpa(q, k, v, mask)
-        B, H, L, D = o.shape
-        return self.out(o.transpose(1, 2).reshape(B, L, H * D))
-
-
-class Block(nn.Module):
-    """Pre-norm residual block: attention of x on a key/value sequence, then an MLP."""
-
-    def __init__(self, dim, n_heads, ff_factor=2, scaling=False, ff_hidden=None):
-        super().__init__()
-        self.norm_q = nn.RMSNorm(dim)
-        self.norm_kv = nn.RMSNorm(dim)
-        self.norm_ff = nn.RMSNorm(dim)
-        self.attn = Attention(dim, n_heads, scaling)
-        self.mlp = MLP(dim, dim * ff_factor if ff_hidden is None else ff_hidden)
-
-    def keys_values(self, kv_input):
-        return self.attn.keys_values(self.norm_kv(kv_input))
-
-    def forward(self, x, k, v, mask=None, q_rope=None, kv_heads=None):
-        x = x + self.attn.attend(self.norm_q(x), k, v, mask, q_rope, kv_heads)
-        return x + self.mlp(self.norm_ff(x))
-
-
-class FoldedRMSNorm(nn.Module):
-    """RMSNorm whose weight was folded into the linear layer that reads it (inference only): one
-    reduction and one multiply instead of the unfused CPU kernels. Same eps as nn.RMSNorm (eps=None:
-    that of the accumulation dtype, float32 for reduced precision), reduction in at least float32."""
-
-    def __init__(self, dim, eps=None):
-        super().__init__()
-        self.dim, self.eps = dim, eps
-
-    def forward(self, x):
-        acc = torch.promote_types(x.dtype, torch.float32)
-        eps = torch.finfo(acc).eps if self.eps is None else self.eps
-        ms = torch.linalg.vector_norm(x, dim=-1, keepdim=True, dtype=acc).square_().div_(self.dim)
-        return x * ms.add_(eps).rsqrt_().to(x.dtype)
-
-
-@torch.no_grad()
-def fold_norm(norm, *linears):
-    """Moves the weight of an RMSNorm into the linear layers fed by it; returns the weightless norm."""
-    for lin in linears:
-        lin.weight.mul_(norm.weight)
-    return FoldedRMSNorm(norm.weight.numel(), norm.eps)
-
-
-@torch.no_grad()
-def interleave_rotary(block):
-    """Reorders the query and key dims of each head of a block from (i, i + D/2) pairs to adjacent
-    (2i, 2i + 1) pairs: q.k is unchanged (same permutation on both), and rope(..., interleaved=True)
-    rotates the same pairs."""
-    a = block.attn
-    H, D = a.n_heads, a.head_dim
-    pairs = torch.stack([torch.arange(D // 2), torch.arange(D // 2) + D // 2], -1).flatten()
-    rows = (torch.arange(H)[:, None] * D + pairs).flatten().to(a.q.weight.device)
-    a.q.weight.copy_(a.q.weight[rows])
-    a.kv.weight[: H * D].copy_(a.kv.weight[rows])
-
-
-class OrthogonalEmbedding(nn.Embedding):
-    """Label embedding with orthonormal initial rows (TabPFN-3 TrainableOrthogonalEmbedding)."""
-
-    def __init__(self, n, dim):
-        super().__init__(n, dim)
-        with torch.no_grad():
-            q, _ = torch.linalg.qr(torch.randn(dim, min(n, dim)))
-            self.weight[: q.shape[1]] = q.T

package/lightpfn/model/lightpfn.py

@@ -1,490 +0,0 @@
-"""LightPFN classifier: cell embedding -> column stage -> row stage -> in-context learning -> decoder.
-
-Design notes and sources are in docs/design.md. In short:
-  cells    per-column statistics of the training rows turn each value into a z-score, an ECDF rank
-           and a NaN flag; cells are grouped with circular feature shifts and embedded with learnable
-           Fourier features (TabPFN-3.5), which avoids the low-rank collapse of a scalar linear
-           embedding (LimiX-2M)
-  column   per-column ISAB whose inducing points attend to the training cells only, with the label
-           embedding added to the training cells (target-aware, TabICLv2)
-  row      per-row transformer over the features with CLS tokens and RoPE; the CLS outputs are
-           concatenated into the row representation (TabICLv2)
-  ICL      transformer over rows: training rows (plus learned thinking rows) attend to each other,
-           test rows attend to the training rows only, with multi-query attention (TabPFN-3.5) and
-           learned log-n softmax scaling (TabPFN-3)
-  decoder  attention of test rows over the one-hot training labels (TabPFN-3 retrieval decoder)
-
-`forward` runs `encode` (everything that depends on the training rows, returned as a Context) and
-then `predict_logits` (test rows only), so the cached inference path is the training path.
-"""
-
-import copy
-from dataclasses import dataclass, field
-
-import torch
-import torch.nn.functional as F
-from torch import nn
-
-from lightpfn.model.layers import Block, OrthogonalEmbedding, SoftmaxScaling, fold_norm, interleave_rotary, rope
-
-
-@dataclass
-class Config:
-    """Opt-in architecture ablations; default parameters/state names are r2-exact.
-
-    Budget recipes at the default widths (trainable parameters, buffers excluded):
-      B1 cell_embed='rbf', icl_ff_reallocation=4: 4,777,024.
-      B2 ccmm=True, icl_ff_reallocation=5: 4,776,784.
-      B3 row_mode='summary': 4,777,072 (no extra parameters).
-      B4 row_refine=True, row_mode='summary', icl_drop_blocks=1: 4,603,088.
-    One last-MLP hidden unit costs 2*icl_dim parameters (512 at default width).
-    B4 uses R=3 independent rounds plus final broadcast: +372,032 parameters,
-    funded by removing one 546,016-parameter ICL block (8 -> 7). Other depths/
-    widths require an explicit budget choice; there is no automatic reallocation.
-    In particular, row_refine does not implicitly change row_mode or ICL depth.
-    """
-
-    col_dim: int = 64
-    col_blocks: int = 2
-    col_heads: int = 4
-    n_inducing: int = 64
-    row_blocks: int = 3
-    row_heads: int = 4
-    n_cls: int = 4  # ICL width = n_cls * col_dim
-    icl_blocks: int = 8
-    icl_heads: int = 8
-    icl_kv_heads_test: int = 1
-    n_thinking: int = 16
-    ff_factor: int = 2
-    n_freq: int = 16
-    n_ecdf_freq: int = 4
-    group_offsets: tuple = (0, 1, 3)
-    label_slots: int = 16  # max classes; training maps classes to random slots so all are trained
-    decoder_heads: int = 4
-    rope_base: float = 10000.0
-    cell_embed: str = "fourier"  # B1: "rbf", 64 fixed uniform kernels, sigma=1
-    ccmm: bool = False  # B2: training-only mask token and rank-bin readout
-    ccmm_mask_fraction: float = 0.15  # observed test cells, per task; structured masks round up
-    row_mode: str = "self_attention"  # B3: "summary"
-    row_refine: bool = False  # B4: refinement -> independent column stage -> compression
-    row_refine_rounds: int = 3
-    icl_drop_blocks: int = 0  # explicit budget reallocation, e.g. 1 for B4
-    icl_ff_reallocation: int = 0  # hidden units removed from the LAST retained ICL MLP
-
-    def __post_init__(self):
-        if self.cell_embed not in ("fourier", "rbf"):
-            raise ValueError("cell_embed must be fourier or rbf")
-        if self.row_mode not in ("self_attention", "summary"):
-            raise ValueError("row_mode must be self_attention or summary")
-        if not 0 <= self.ccmm_mask_fraction <= 1:
-            raise ValueError("ccmm_mask_fraction must be in [0, 1]")
-        if self.row_refine_rounds < 1:
-            raise ValueError("row_refine_rounds must be positive")
-        if not 0 <= self.icl_drop_blocks < self.icl_blocks:
-            raise ValueError("icl_drop_blocks must leave at least one ICL block")
-        if not 0 <= self.icl_ff_reallocation < self.icl_dim * self.ff_factor:
-            raise ValueError("icl_ff_reallocation must leave a nonempty ICL MLP")
-
-    @property
-    def icl_dim(self):
-        return self.n_cls * self.col_dim
-
-
-@dataclass
-class Context:
-    """Everything prediction needs from the training rows."""
-
-    stats: dict
-    col_kv: list
-    icl_kv: list
-    dec_k: torch.Tensor
-    y: torch.Tensor
-    slots: torch.Tensor
-    d: torch.Tensor | None
-    n_classes: int
-    extra: dict = field(default_factory=dict)
-    col_refine_kv: list | None = None
-
-
-class CellEmbedder(nn.Module):
-    def __init__(self, cfg):
-        super().__init__()
-        G = len(cfg.group_offsets)
-        self.offsets = cfg.group_offsets
-        self.n_ecdf_freq = cfg.n_ecdf_freq
-        self.cell_embed = cfg.cell_embed
-        if cfg.cell_embed == "fourier":
-            self.freq = nn.Parameter(torch.randn(G, cfg.n_freq) * 2.0)
-            self.fourier = nn.Linear(2 * cfg.n_freq, cfg.col_dim, bias=False)
-        else:
-            # RaBEL's verified sweep favors 64 uniform kernels and fixed sigma=1.
-            # The range is adapted to LightPFN's already soft-clipped z scores.
-            self.register_buffer("rbf_centers", torch.linspace(-5.0, 5.0, 64))
-            self.rbf = nn.Linear(64, cfg.col_dim, bias=False)
-        self.meta = nn.Linear(G * (2 + 2 * cfg.n_ecdf_freq), cfg.col_dim, bias=False)
-        self.norm = nn.LayerNorm(cfg.col_dim)
-
-    @staticmethod
-    @torch.no_grad()
-    def stats(x_train):
-        """Column statistics of the training rows (NaN ignored). x_train: (B, n, m)."""
-        x = x_train.float()
-        valid = ~torch.isnan(x)
-        cnt = valid.sum(1)  # (B, m)
-        x0 = torch.where(valid, x, 0.0)
-        mean = x0.sum(1) / cnt.clamp(min=1)
-        var = (torch.where(valid, x - mean[:, None], 0.0) ** 2).sum(1) / (cnt - 1).clamp(min=1)
-        std = var.sqrt()
-        srt = torch.where(valid, x, float("inf")).transpose(1, 2).sort(-1).values.contiguous()  # (B, m, n)
-        return dict(mean=mean, std=std, sorted=srt, cnt=cnt)
-
-    @staticmethod
-    def normalize(x, st):
-        """Returns z-score (soft-clipped), ECDF mid-rank in [0, 1] and NaN flag, each (B, n, m)."""
-        x = x.float()
-        nan = torch.isnan(x)
-        z = (x - st["mean"][:, None]) / (st["std"][:, None] + 1e-6)
-        z = torch.where(nan | (st["std"][:, None] == 0), 0.0, 5.0 * torch.tanh(z / 5.0))
-        q = torch.where(nan, 0.0, x).transpose(1, 2).contiguous()  # (B, m, n)
-        lo = torch.searchsorted(st["sorted"], q, right=False)
-        hi = torch.searchsorted(st["sorted"], q, right=True)
-        cnt = st["cnt"][..., None]
-        r = ((lo + hi).float() / 2 / cnt.clamp(min=1)).clamp(0, 1)
-        r = torch.where(cnt > 0, r, 0.5).transpose(1, 2)
-        r = torch.where(nan, 0.5, r)
-        return z, r, nan.float()
-
-    def group_index(self, d, B, m, device):
-        j = torch.arange(m, device=device)[None]
-        d = torch.full((B, 1), m, device=device) if d is None else d.to(device)[:, None]
-        return torch.stack([torch.where(j < d, (j + o) % d, j) for o in self.offsets], -1)  # (B, m, G)
-
-    def grouped(self, x, st, d=None, mask=None, mask_token=None):
-        """z, ECDF rank and NaN flag of every cell and of its group neighbors, each (B, n, m, G)."""
-        B, n, m = x.shape
-        if mask is None:
-            z, r, nan = self.normalize(x, st)
-        else:
-            if mask.shape != x.shape or mask.dtype != torch.bool or mask_token is None:
-                raise ValueError("cell masking needs a bool mask matching x and a mask token")
-            # Neutralize BEFORE circular grouping: otherwise a hidden value leaks into
-            # the embeddings of its neighbors through the grouped z/ECDF/NaN features.
-            z, r, nan = self.normalize(x.masked_fill(mask, 0.0), st)
-            z, r, nan = (t.masked_fill(mask, 0.0) for t in (z, r, nan))
-        idx = self.group_index(d, B, m, x.device)
-        G = idx.shape[-1]
-        gidx = idx.view(B, 1, m * G).expand(B, n, m * G)
-        return tuple(torch.gather(t, 2, gidx).view(B, n, m, G) for t in (z, r, nan))
-
-    def forward(self, x, st, d=None, mask=None, mask_token=None):
-        out = self.embed(*self.grouped(x, st, d, mask, mask_token))
-        if mask is not None:
-            out = torch.where(mask[..., None], mask_token.to(out.dtype), out)
-        return out
-
-    def embed(self, z, r, nan):
-        """Cell embeddings (B, n, m, E) from grouped features; cells are independent, so any slice
-        of rows or columns of the features gives the same slice of the embeddings."""
-        if self.cell_embed == "fourier":
-            ang = z[..., None] * self.freq  # (B, n, m, G, F), fp32
-            four = torch.cat([ang.sin(), ang.cos()], -1).sum(-2)  # (B, n, m, 2F)
-            value = self.fourier(four)
-        else:
-            kernels = torch.exp(-0.5 * (z[..., None] - self.rbf_centers).square())
-            value = self.rbf(kernels.sum(-2))
-        k = torch.pi * 2.0 ** torch.arange(self.n_ecdf_freq, device=z.device)
-        rang = r[..., None] * k
-        meta = torch.cat([z, nan, rang.sin().flatten(-2), rang.cos().flatten(-2)], -1)
-        return self.norm(value + self.meta(meta))  # (B, n, m, E)
-
-
-class ColumnStage(nn.Module):
-    """ISAB over the rows of each column. Inducing points read the training cells only, so their
-    states (col_kv) are all that test rows need."""
-
-    def __init__(self, cfg):
-        super().__init__()
-        E = cfg.col_dim
-        self.y_emb = OrthogonalEmbedding(cfg.label_slots, E)
-        self.inducing = nn.ParameterList(nn.Parameter(torch.randn(cfg.n_inducing, E) * 0.02) for _ in range(cfg.col_blocks))
-        self.ind_blocks = nn.ModuleList(Block(E, cfg.col_heads, cfg.ff_factor) for _ in range(cfg.col_blocks))
-        self.cell_blocks = nn.ModuleList(Block(E, cfg.col_heads, cfg.ff_factor) for _ in range(cfg.col_blocks))
-
-    def context(self, cells, y_slots):
-        B, n, m, E = cells.shape
-        x = (cells + self.y_emb(y_slots)[:, :, None]).transpose(1, 2).reshape(B * m, n, E)
-        col_kv = []
-        for ind, ib, cb in zip(self.inducing, self.ind_blocks, self.cell_blocks):
-            h = ib(ind.expand(B * m, -1, -1), *ib.keys_values(x))
-            kh, vh = cb.keys_values(h)
-            x = cb(x, kh, vh)
-            col_kv.append((kh, vh))
-        return x.view(B, m, n, E).transpose(1, 2), col_kv
-
-    def query(self, cells, col_kv):
-        B, n, m, E = cells.shape
-        x = cells.transpose(1, 2).reshape(B * m, n, E)
-        for cb, (kh, vh) in zip(self.cell_blocks, col_kv):
-            x = cb(x, kh, vh)
-        return x.view(B, m, n, E).transpose(1, 2)
-
-
-class RowStage(nn.Module):
-    """Transformer over the features of each row; CLS tokens (no rotation) collect the row."""
-
-    def __init__(self, cfg):
-        super().__init__()
-        self.n_cls, self.rope_base = cfg.n_cls, cfg.rope_base
-        self.rope_interleaved = False  # set by LightPFN.folded()
-        self.mode = cfg.row_mode
-        self.cls = nn.Parameter(torch.randn(cfg.n_cls, cfg.col_dim) * 0.02)
-        self.blocks = nn.ModuleList(Block(cfg.col_dim, cfg.row_heads, cfg.ff_factor) for _ in range(cfg.row_blocks))
-
-    def _rope(self, t):
-        C = self.n_cls
-        pos = torch.arange(t.shape[2] - C, device=t.device)
-        return torch.cat([t[:, :, :C], rope(t[:, :, C:], pos, self.rope_base, self.rope_interleaved)], dim=2)
-
-    def forward(self, cells, d=None, has_padding=None, return_cells=False):
-        B, n, m, E = cells.shape
-        C = self.n_cls
-        x = torch.cat([self.cls.to(cells.dtype).expand(B * n, C, E), cells.reshape(B * n, m, E)], dim=1)
-        mask = None
-        # Training supplies the exact decision from CPU metadata. Other callers
-        # retain the original automatic padding detection.
-        if d is not None and (bool((d < m).any()) if has_padding is None else has_padding):
-            valid = torch.arange(m, device=cells.device)[None] < d.to(cells.device)[:, None]  # (B, m)
-            valid = torch.cat([torch.ones(B, C, dtype=torch.bool, device=cells.device), valid], 1)
-            mask = valid.repeat_interleave(n, 0)[:, None, None, :]  # (B*n, 1, 1, C+m)
-        for blk in self.blocks:
-            if self.mode == "self_attention":
-                k, v = blk.keys_values(x)
-                x = blk(x, self._rope(k), v, mask, q_rope=self._rope)
-            else:
-                # Only C queries: cells stay fixed; summaries also attend to each other.
-                k, v = blk.keys_values(x)
-                summary = blk(x[:, :C], self._rope(k), v, mask)
-                x = torch.cat([summary, x[:, C:]], dim=1)
-        rows = x[:, :C].reshape(B, n, C * E)
-        if return_cells:
-            return rows, x[:, C:].reshape(B, n, m, E)
-        return rows
-
-
-class RowRefinement(nn.Module):
-    """Temporary summaries, independent broadcast/gather weights per round, final broadcast.
-
-    Each row is processed independently. Feature attention is O(m*K), K=4;
-    summaries are discarded before the second column stage and final compression.
-    """
-
-    def __init__(self, cfg):
-        super().__init__()
-        self.rope_base = cfg.rope_base
-        self.rope_interleaved = False  # set by LightPFN.folded()
-        self.summary = nn.Parameter(torch.randn(4, cfg.col_dim) * 0.02)
-        self.broadcast = nn.ModuleList(
-            Block(cfg.col_dim, cfg.row_heads, cfg.ff_factor) for _ in range(cfg.row_refine_rounds + 1))
-        self.gather = nn.ModuleList(
-            Block(cfg.col_dim, cfg.row_heads, cfg.ff_factor) for _ in range(cfg.row_refine_rounds))
-
-    def forward(self, cells, d=None, has_padding=None):
-        B, n, m, E = cells.shape
-        x = cells.reshape(B * n, m, E)
-        s = self.summary.to(cells.dtype).expand(B * n, -1, -1)
-        mask = None
-        if d is not None and (bool((d < m).any()) if has_padding is None else has_padding):
-            valid = torch.arange(m, device=cells.device)[None] < d.to(cells.device)[:, None]
-            mask = valid.repeat_interleave(n, 0)[:, None, None, :]
-        pos = torch.arange(m, device=cells.device)
-        for i, broadcast in enumerate(self.broadcast):
-            x = broadcast(x, *broadcast.keys_values(s),
-                          q_rope=lambda q: rope(q, pos, self.rope_base, self.rope_interleaved))
-            if i < len(self.gather):
-                gather = self.gather[i]
-                k, v = gather.keys_values(x)
-                s = gather(s, rope(k, pos, self.rope_base, self.rope_interleaved), v, mask)
-        return x.reshape(B, n, m, E)
-
-
-class ICLStage(nn.Module):
-    def __init__(self, cfg):
-        super().__init__()
-        D = cfg.icl_dim
-        self.kv_heads_test = cfg.icl_kv_heads_test
-        self.y_emb = OrthogonalEmbedding(cfg.label_slots, D)
-        self.thinking = nn.Parameter(torch.randn(cfg.n_thinking, D) * 0.02)
-        depth = cfg.icl_blocks - cfg.icl_drop_blocks
-        self.blocks = nn.ModuleList(
-            Block(D, cfg.icl_heads, cfg.ff_factor, scaling=True,
-                  ff_hidden=D * cfg.ff_factor - (cfg.icl_ff_reallocation if i == depth - 1 else 0))
-            for i in range(depth))
-        self.norm = nn.RMSNorm(D)
-
-    def context(self, r, y_slots):
-        B = r.shape[0]
-        T = self.thinking.shape[0]
-        x = torch.cat([self.thinking.to(r.dtype).expand(B, -1, -1), r + self.y_emb(y_slots)], dim=1)
-        icl_kv = []
-        for blk in self.blocks:
-            k, v = blk.keys_values(x)
-            x = blk(x, k, v)
-            if not self.training:  # the cache keeps only the heads test rows read
-                k, v = k[:, : self.kv_heads_test].contiguous(), v[:, : self.kv_heads_test].contiguous()
-            icl_kv.append((k, v))
-        return self.norm(x[:, T:]), icl_kv
-
-    def query(self, r, icl_kv):
-        x = r
-        for blk, (k, v) in zip(self.blocks, icl_kv):
-            x = blk(x, k, v, kv_heads=self.kv_heads_test)
-        return self.norm(x)
-
-
-class RetrievalDecoder(nn.Module):
-    """p(class | test row) = attention-weighted average of the one-hot training labels, averaged
-    over heads; logits are its log. Any number of classes up to the head dim."""
-
-    def __init__(self, cfg):
-        super().__init__()
-        D, H = cfg.icl_dim, cfg.decoder_heads
-        self.n_heads, self.head_dim = H, D // H
-        assert cfg.label_slots <= self.head_dim
-        self.q = nn.Linear(D, D, bias=False)
-        self.k = nn.Linear(D, D, bias=False)
-        self.scaling = SoftmaxScaling(H, self.head_dim)
-
-    def keys(self, h_train):
-        B, N, _ = h_train.shape
-        return self.k(h_train).view(B, N, self.n_heads, self.head_dim).transpose(1, 2)
-
-    def forward(self, k, y, h_test, n_classes):
-        B, M, _ = h_test.shape
-        q = self.q(h_test).view(B, M, self.n_heads, self.head_dim).transpose(1, 2)
-        q = self.scaling(q, k.shape[2])
-        # one-hot values padded to the head dim so fused attention kernels apply
-        v = F.one_hot(y, self.head_dim).to(q.dtype)[:, None].expand(-1, self.n_heads, -1, -1)
-        p = F.scaled_dot_product_attention(q, k.to(q.dtype), v).float().mean(1)[..., :n_classes]
-        return torch.log(p.clamp(min=1e-5) + 3e-5)
-
-
-class LightPFN(nn.Module):
-    def __init__(self, cfg=None):
-        super().__init__()
-        self.cfg = cfg = cfg or Config()
-        self.cells = CellEmbedder(cfg)
-        self.col = ColumnStage(cfg)
-        self.row = RowStage(cfg)
-        self.icl = ICLStage(cfg)
-        self.decoder = RetrievalDecoder(cfg)
-        if cfg.row_refine:
-            self.refine = RowRefinement(cfg)
-            self.col_refine = ColumnStage(cfg)  # independent, target-aware, train-only cache
-        if cfg.ccmm:
-            self.ccmm_mask_token = nn.Parameter(torch.randn(cfg.col_dim) * 0.02)
-            self.ccmm_head = nn.Sequential(nn.LayerNorm(cfg.col_dim), nn.Linear(cfg.col_dim, 32))
-
-    @torch.no_grad()
-    def folded(self):
-        """A copy for prediction only that computes the same function faster: RMSNorm weights folded
-        into the linear layers that read them, and the query/key dims of the rotary blocks interleaved
-        so RoPE is one complex multiply. Outputs match the original up to float rounding; the copy is
-        not meant for training or for saving as a checkpoint."""
-        if getattr(self, "is_folded", False):
-            return self
-        m = copy.deepcopy(self).eval()
-        m.is_folded = True
-        for blk in (b for b in m.modules() if isinstance(b, Block)):
-            blk.norm_q = fold_norm(blk.norm_q, blk.attn.q)
-            blk.norm_kv = fold_norm(blk.norm_kv, blk.attn.kv)
-            blk.norm_ff = fold_norm(blk.norm_ff, blk.mlp.fc1)
-        m.icl.norm = fold_norm(m.icl.norm, m.decoder.q, m.decoder.k)  # its output only feeds the decoder
-        rotary = [m.row] + ([m.refine] if self.cfg.row_refine else [])
-        for stage in rotary:
-            for blk in (b for b in stage.modules() if isinstance(b, Block)):
-                interleave_rotary(blk)
-            stage.rope_interleaved = True
-        return m
-
-    def encode(self, X_train, y_train, d=None, slots=None, n_classes=None, has_padding=None, chunk_cells=None):
-        """X_train: (B, n, m) float with NaN, y_train: (B, n) labels 0..C-1, d: (B,) true feature
-        counts when features are zero-padded, slots: (B, label_slots) class -> label slot map.
-        n_classes: total number of classes C. The default, max(y_train) + 1, is wrong when the
-        highest classes have no training row: callers that know C must pass it (the sklearn
-        wrapper and the training loop do). chunk_cells (inference): run the cell stages on about
-        that many cells at a time, which keeps them in the CPU cache; the result is the same."""
-        B = X_train.shape[0]
-        n_classes = int(y_train.max()) + 1 if n_classes is None else n_classes
-        if n_classes > self.cfg.label_slots:
-            raise ValueError(f"{n_classes} classes, the model supports {self.cfg.label_slots}")
-        if slots is None:
-            slots = torch.arange(self.cfg.label_slots, device=X_train.device).expand(B, -1)
-        y_slots = torch.gather(slots, 1, y_train)
-        stats = self.cells.stats(X_train)
-        if chunk_cells is not None:
-            if torch.is_grad_enabled():
-                raise RuntimeError("chunk_cells is for inference: it writes into a shared buffer that autograd cannot track")
-            rows, col_kv, col_refine_kv = self._train_rows_chunked(X_train, stats, y_slots, d, has_padding, chunk_cells)
-        else:
-            col, col_kv = self.col.context(self.cells(X_train, stats, d), y_slots)
-            col_refine_kv = None
-            if self.cfg.row_refine:
-                col = self.refine(col, d, has_padding)
-                col, col_refine_kv = self.col_refine.context(col, y_slots)
-            rows = self.row(col, d, has_padding)
-        h, icl_kv = self.icl.context(rows, y_slots)
-        return Context(stats, col_kv, icl_kv, self.decoder.keys(h), y_train, slots, d, n_classes,
-                       extra=dict(has_padding=has_padding), col_refine_kv=col_refine_kv)
-
-    def _train_rows_chunked(self, X, stats, y_slots, d, has_padding, chunk_cells):
-        """The cell stages of encode by pieces: column stages on groups of whole columns, row stages
-        on groups of whole rows, all writing into one (B, n, m, E) buffer."""
-        B, n, m = X.shape
-        grouped = self.cells.grouped(X, stats, d)
-        cols, rows = max(1, chunk_cells // (B * n)), max(1, chunk_cells // (B * m))
-        buf = None  # (B, n, m, E), allocated with the dtype of the first column-stage output
-
-        def column_stage(stage, cells_of):
-            nonlocal buf
-            parts = []
-            for j in range(0, m, cols):
-                out, kv = stage.context(cells_of(j, j + cols), y_slots)
-                if buf is None:
-                    buf = out.new_empty(B, n, m, out.shape[-1])
-                buf[:, :, j : j + cols] = out
-                parts.append(kv)
-            # (B * columns, ...) caches of each piece, merged in the (batch, column) order of one call
-            return [tuple(torch.cat([p[i][t].unflatten(0, (B, -1)) for p in parts], 1).flatten(0, 1) for t in (0, 1))
-                    for i in range(len(parts[0]))]
-
-        col_kv = column_stage(self.col, lambda a, b: self.cells.embed(*(t[:, :, a:b] for t in grouped)))
-        col_refine_kv = None
-        if self.cfg.row_refine:
-            for i in range(0, n, rows):
-                buf[:, i : i + rows] = self.refine(buf[:, i : i + rows], d, has_padding)
-            col_refine_kv = column_stage(self.col_refine, lambda a, b: buf[:, :, a:b])
-        out = torch.cat([self.row(buf[:, i : i + rows], d, has_padding) for i in range(0, n, rows)], 1)
-        return out, col_kv, col_refine_kv
-
-    def _test_rows(self, ctx, X_test):
-        col = self.col.query(self.cells(X_test, ctx.stats, ctx.d), ctx.col_kv)
-        if self.cfg.row_refine:
-            col = self.refine(col, ctx.d, ctx.extra.get("has_padding"))
-            col = self.col_refine.query(col, ctx.col_refine_kv)
-        return self.row(col, ctx.d, ctx.extra.get("has_padding"))
-
-    def predict_logits(self, ctx, X_test, chunk_cells=None):
-        """Logits (B, M, n_classes) for test rows; rows are independent, so X_test can be chunked.
-        chunk_cells: run the cell stages on groups of about that many cells (same result)."""
-        if chunk_cells is None:
-            rows = self._test_rows(ctx, X_test)
-        else:
-            step = max(1, chunk_cells // (X_test.shape[0] * X_test.shape[2]))
-            rows = torch.cat([self._test_rows(ctx, X_test[:, i : i + step]) for i in range(0, X_test.shape[1], step)], 1)
-        h = self.icl.query(rows, ctx.icl_kv)
-        return self.decoder(ctx.dec_k, ctx.y, h, ctx.n_classes)
-
-    def forward(self, X, y_train, d=None, slots=None, n_classes=None, has_padding=None):
-        n_train = y_train.shape[1]
-        ctx = self.encode(X[:, :n_train], y_train, d, slots, n_classes, has_padding)
-        return self.predict_logits(ctx, X[:, n_train:])

package/lightpfn/pretrained.json

@@ -1,4 +0,0 @@
-{
-  "repo_id": "ueuegio/LightPFN",
-  "revision": "bd389ab59a89dd0e05c9ecb7c642c08ee52e9637"
-}

package/lightpfn/sklearn.py

@@ -1,280 +0,0 @@
-"""scikit-learn style wrapper: fit() encodes the training set once (column statistics, inducing
-states, ICL key/value cache), predict_proba() only runs the test rows, in chunks.
-
-Estimators beyond the first use a random feature order and a random class -> label-slot map;
-their probabilities are averaged. Training sets larger than max_context are subsampled per
-estimator with stratification, so every class (even one with a single row) stays in the context.
-"""
-
-import copy
-import warnings
-from numbers import Integral
-
-import numpy as np
-import torch
-from sklearn.base import BaseEstimator, ClassifierMixin
-from sklearn.utils.multiclass import check_classification_targets
-from sklearn.utils.validation import check_is_fitted, validate_data
-
-from lightpfn.checkpoint import load_model, load_pretrained
-from lightpfn.device import kind, resolve_device
-from lightpfn.model.lightpfn import Config, LightPFN
-
-# Inference pieces, in cells (rows x features, times the estimators of a batch); measured on an 8-core
-# CPU and an RTX 5090 (runs/bench_infer/).
-CPU_CHUNK_CELLS_PER_THREAD = 1024
-GPU_CHUNK_CELLS = 1 << 20
-CPU_BATCH_CELLS = 0  # one estimator at a time: batching saves nothing once pieces fit the cache
-GPU_BATCH_CELLS = 1 << 22
-
-
-def stratified_subsample(rng, y, n, min_per_class=5):
-    """Exactly n of the len(y) > n rows: class proportions kept, and min(count, min_per_class) rows
-    of every class, taken from the largest classes (with fewer classes guaranteed if the minimums
-    alone would exceed n)."""
-    classes, counts = np.unique(y, return_counts=True)
-    if n < len(classes):
-        raise ValueError("max_context must be at least the number of classes.")
-    mins = np.minimum(counts, min_per_class)
-    if mins.sum() > n:
-        mins = np.minimum(counts, max(1, n // len(classes)))
-    take = np.maximum(np.floor(counts * n / len(y)).astype(int), mins)
-    while take.sum() > n:
-        take[np.argmax(take - mins)] -= 1
-    idx = np.concatenate([rng.choice(np.flatnonzero(y == c), size=k, replace=False) for c, k in zip(classes, take)])
-    if len(idx) < n:
-        rest = np.setdiff1d(np.arange(len(y)), idx)
-        idx = np.concatenate([idx, rng.choice(rest, size=n - len(idx), replace=False)])
-    return np.sort(idx)
-
-
-class LightPFNClassifier(ClassifierMixin, BaseEstimator):
-    """device: "auto" (default: CUDA/ROCm through torch, else a Vulkan GPU of any vendor, else the CPU; the
-    environment variable LIGHTPFN_DEVICE can replace it), or explicitly "cuda[:i]", "vulkan[:i]", "cpu".
-    On "vulkan" the network runs on lightpfn.vulkan (wgpu); a training set too large for the GPU's buffers
-    falls back to the CPU with a warning when the device was chosen by "auto", and raises otherwise.
-    Weights and devices are initialized on fit, so construction and sklearn clone
-    perform no downloads or GPU work. Without model/checkpoint the pinned release
-    is downloaded from Hugging Face (install the hf extra and authenticate while
-    the repository is private). Input must be numeric; NaN is supported. Encode
-    strings/categories before fit. The model was trained for 2-10 classes.
-
-    n_estimators averages feature/label permutations; max_context bounds the
-    stratified context per member. chunk_rows bounds prediction batches;
-    chunk_cells/batch_cells control cache blocking and estimator batching.
-    n_threads sets PyTorch's process-wide CPU thread count. seed is the original
-    RNG parameter; random_state is its sklearn alias (use one of them).
-    fold enables equivalent inference folding. A supplied torch model is copied
-    during fit, leaving the constructor parameter and its device unchanged.
-    """
-
-    def __init__(self, model=None, checkpoint=None, device="auto", n_estimators=1, max_context=20000,
-                 chunk_rows=2048, n_threads=None, seed=0, *, chunk_cells="auto", batch_cells="auto", fold=True,
-                 random_state=None, repo_id=None, revision=None, cache_dir=None, local_files_only=False):
-        self.model = model
-        self.checkpoint = checkpoint
-        self.device = device
-        self.fold = fold
-        self.n_estimators = n_estimators
-        self.max_context = max_context
-        self.chunk_rows = chunk_rows
-        self.chunk_cells = chunk_cells
-        self.batch_cells = batch_cells
-        self.n_threads = n_threads
-        self.seed = seed
-        self.random_state = random_state
-        self.repo_id = repo_id
-        self.revision = revision
-        self.cache_dir = cache_dir
-        self.local_files_only = local_files_only
-
-    def __sklearn_tags__(self):
-        tags = super().__sklearn_tags__()
-        tags.input_tags.allow_nan = True
-        # A pretrained network is not optimized on sklearn's toy check datasets.
-        tags.classifier_tags.poor_score = True
-        return tags
-
-    def _initialize_backend(self):
-        key = (id(self.model), self.checkpoint, self.device, self.fold, self.repo_id,
-               self.revision, self.cache_dir, self.local_files_only)
-        if getattr(self, "_backend_key", None) == key:
-            return
-        self.device_ = resolve_device(self.device)
-        torch_device = "cpu" if kind(self.device_) == "vulkan" else self.device_
-        if self.model is not None and self.checkpoint is not None:
-            raise ValueError("Provide either model or checkpoint, not both.")
-        if self.model is not None:
-            self.model_ = copy.deepcopy(self.model) if isinstance(self.model, torch.nn.Module) else self.model
-        elif self.checkpoint is not None:
-            self.model_ = load_model(self.checkpoint, torch_device)
-        else:
-            self.model_ = load_pretrained(repo_id=self.repo_id, revision=self.revision, device=torch_device,
-                                         cache_dir=self.cache_dir, local_files_only=self.local_files_only)
-        if isinstance(self.model_, torch.nn.Module):
-            self.model_ = self.model_.to(torch_device).eval()
-        # fold=True predicts with LightPFN.folded(), the same function with fewer memory passes
-        self.cpu_net_ = None
-        if kind(self.device_) == "vulkan":
-            from lightpfn.vulkan import VulkanLightPFN
-
-            index = self.device_.partition(":")[2]
-            self.net_ = VulkanLightPFN(self.model_, adapter=int(index) if index else None)
-        else:
-            self.net_ = self.model_.folded() if self.fold else self.model_
-        self._backend_key = key
-
-    def _validate_parameters(self):
-        for name in ("n_estimators", "max_context", "chunk_rows"):
-            value = getattr(self, name)
-            if isinstance(value, bool) or not isinstance(value, Integral) or value < 1:
-                raise ValueError(f"{name} must be a positive integer.")
-        if self.n_threads is not None and (isinstance(self.n_threads, bool) or
-                not isinstance(self.n_threads, Integral) or self.n_threads < 1):
-            raise ValueError("n_threads must be None or a positive integer.")
-        for name in ("chunk_cells", "batch_cells"):
-            value = getattr(self, name)
-            if value == "auto" or (name == "chunk_cells" and value is None):
-                continue
-            minimum = 1 if name == "chunk_cells" else 0
-            if isinstance(value, bool) or not isinstance(value, Integral) or value < minimum:
-                raise ValueError(f"{name} must be 'auto' or an integer >= {minimum}." )
-        if not isinstance(self.fold, bool):
-            raise ValueError("fold must be a boolean.")
-        if self.random_state is not None and self.seed not in (0, None):
-            raise ValueError("Use random_state or seed, not both.")
-
-    def __sklearn_is_fitted__(self):
-        return getattr(self, "_is_fitted", False)
-
-    def _auto(self, value, cpu, gpu):
-        if value != "auto":
-            return value
-        return cpu if kind(self.fit_device_) == "cpu" else gpu
-
-    def _chunk_cells(self):
-        """Cells per piece of the cell stages: small pieces stay in the CPU cache (about 3x faster on
-        wide tables), large pieces bound GPU memory. The predictions do not depend on it."""
-        return self._auto(self.chunk_cells, CPU_CHUNK_CELLS_PER_THREAD * torch.get_num_threads(), GPU_CHUNK_CELLS)
-
-    def _group(self, n, m):
-        """Estimators encoded together as one batch: as many as fit in batch_cells training cells
-        (fewer, larger operations; the same predictions as one at a time)."""
-        return max(1, min(self.n_estimators, self._auto(self.batch_cells, CPU_BATCH_CELLS, GPU_BATCH_CELLS) // max(1, n * m)))
-
-    @torch.inference_mode()
-    def fit(self, X, y, cat=None):
-        self._is_fitted = False
-        self._validate_parameters()
-        X, y = validate_data(self, X, y, dtype=np.float32, ensure_all_finite="allow-nan")
-        check_classification_targets(y)
-        if cat is not None:
-            cat = np.asarray(cat)
-            if cat.shape != (X.shape[1],) or cat.dtype != np.bool_:
-                raise ValueError("cat must be a boolean mask with one entry per feature.")
-            if cat.any():
-                warnings.warn("cat does not enable native categorical handling; columns are treated as numeric "
-                              "codes. Encode categories before fit. The cat parameter is deprecated.",
-                              FutureWarning, stacklevel=2)
-        if self.n_threads is not None:
-            torch.set_num_threads(self.n_threads)
-        self.classes_, y = np.unique(np.asarray(y), return_inverse=True)
-        if len(self.classes_) > 10:
-            raise ValueError("LightPFN supports at most 10 classes (trained on 2-10 classes).")
-        if self.max_context < len(self.classes_):
-            raise ValueError("max_context must be at least the number of classes.")
-        self._initialize_backend()
-        if len(self.classes_) > self.net_.cfg.label_slots:
-            raise ValueError("Number of classes exceeds the model's label slots.")
-        random_state = self.seed if self.random_state is None else self.random_state
-        if isinstance(random_state, np.random.RandomState):
-            random_state = random_state.randint(2**32)
-        rng = np.random.default_rng(random_state)
-        S = self.net_.cfg.label_slots
-        plans = []
-        for e in range(self.n_estimators):
-            rows = np.arange(len(y))
-            if len(rows) > self.max_context:
-                rows = stratified_subsample(rng, y, self.max_context)
-            feats = np.arange(X.shape[1]) if e == 0 else rng.permutation(X.shape[1])
-            slots = np.arange(S) if e == 0 else rng.permutation(S)
-            plans.append((rows, feats, slots))
-        self._choose_backend(len(plans[0][0]), X.shape[1])
-        try:
-            self._encode_members(X, y, plans)
-        except MemoryError as exc:
-            if kind(self.fit_device_) != "vulkan":
-                raise
-            # Drop partial contexts and unsubmitted work before a retry or a later fit.
-            msg = str(exc)
-            exc.__traceback__ = None
-            self.members_.clear()
-            eng = self.fit_net_.eng
-            eng.ops.clear()
-            eng.pending = 0.0
-            eng.scratch.clear()
-            eng.finish()
-            self._use_cpu(msg)
-            self._encode_members(X, y, plans)
-        if kind(self.fit_device_) == "cuda":
-            torch.cuda.synchronize(self.fit_device_)  # fit time includes the GPU work, not only its launch
-        elif kind(self.fit_device_) == "vulkan":
-            self.fit_net_.eng.finish()
-        self._is_fitted = True
-        return self
-
-    def _encode_members(self, X, y, plans):
-        group = self._group(len(plans[0][0]), X.shape[1])
-        if kind(self.fit_device_) == "vulkan":
-            group = min(group, self.fit_net_.train_capacity(len(plans[0][0]), X.shape[1]))
-        tdev = self._tensor_device()
-        self.members_ = []
-        for g in range(0, len(plans), group):
-            part = plans[g : g + group]
-            Xt = torch.from_numpy(np.stack([X[rows][:, feats] for rows, feats, _ in part])).to(tdev)
-            yt = torch.from_numpy(np.stack([y[rows] for rows, _, _ in part])).long().to(tdev)
-            st = torch.from_numpy(np.stack([slots for _, _, slots in part])).long().to(tdev)
-            ctx = self.fit_net_.encode(Xt, yt, slots=st, n_classes=len(self.classes_), chunk_cells=self._chunk_cells())
-            # (feature order, context): one estimator per entry as before, or (G, m) orders for a batch of G
-            feats = part[0][1] if len(part) == 1 else np.stack([feats for _, feats, _ in part])
-            self.members_.append((feats, ctx))
-
-    def _choose_backend(self, n, m):
-        """The network and device of this fit: the classifier's, or the CPU when a Vulkan device chosen by
-        "auto" cannot hold one estimator's caches and scratch buffers."""
-        self.fit_device_, self.fit_net_ = self.device_, self.net_
-        if kind(self.device_) == "vulkan" and self.net_.train_capacity(n, m) < 1:
-            self._use_cpu(f"{n} x {m} training cells exceed a Vulkan cache/scratch buffer limit; "
-                          "lower max_context or use device='cpu'")
-
-    def _use_cpu(self, msg):
-        if not resolve_device_was_auto(self.device):
-            raise MemoryError(msg)
-        warnings.warn(msg + ": running this fit on the CPU", RuntimeWarning, stacklevel=3)
-        if self.cpu_net_ is None:
-            self.cpu_net_ = self.model_.folded() if self.fold else self.model_
-        self.fit_device_, self.fit_net_ = "cpu", self.cpu_net_
-
-    def _tensor_device(self):
-        return "cpu" if kind(self.fit_device_) == "vulkan" else self.fit_device_
-
-    @torch.inference_mode()
-    def predict_proba(self, X):
-        check_is_fitted(self)
-        X = validate_data(self, X, reset=False, dtype=np.float32, ensure_all_finite="allow-nan")
-        P = np.zeros((len(X), len(self.classes_)))
-        for feats, ctx in self.members_:
-            for i in range(0, len(X), self.chunk_rows):
-                Xt = torch.from_numpy(np.stack([X[i : i + self.chunk_rows][:, f] for f in np.atleast_2d(feats)])).to(self._tensor_device())
-                probs = torch.softmax(self.fit_net_.predict_logits(ctx, Xt, self._chunk_cells()).float(), -1).cpu().numpy()
-                for p in probs:  # one estimator at a time, in float64 as before
-                    P[i : i + self.chunk_rows] += p
-        return P / self.n_estimators
-
-    def predict(self, X):
-        probabilities = self.predict_proba(X)
-        return self.classes_[probabilities.argmax(1)]
-
-
-def resolve_device_was_auto(device):
-    return device is None or str(device).strip().lower() == "auto"

package/lightpfn/vulkan/__init__.py

@@ -1,40 +0,0 @@
-"""Vulkan backend: LightPFN inference on any GPU with a Vulkan driver (AMD, Intel, NVIDIA; Linux and
-Windows), through wgpu (`pip install wgpu`). The WGSL kernels in kernels.py are compiled to SPIR-V when the
-backend starts; no Vulkan SDK or compiler is needed.
-
-    from lightpfn.vulkan import VulkanLightPFN, is_available
-    net = VulkanLightPFN(model)              # first GPU with a Vulkan driver
-    ctx = net.encode(X_train, y_train, n_classes=C)
-    logits = net.predict_logits(ctx, X_test)
-
-LIGHTPFN_VULKAN_ADAPTER selects the adapter by index or name ("llvmpipe" runs the kernels on the CPU,
-which is how the tests run without a GPU).
-"""
-
-from lightpfn.vulkan.engine import ADAPTER_ENV, GPU_TYPES, pick_adapter, vulkan_adapters
-
-
-def adapters():
-    """The Vulkan adapters wgpu sees, as dicts (index, name, type, vendor, driver)."""
-    return [dict(index=i, name=info.get("device"), type=info.get("adapter_type"), vendor=info.get("vendor"),
-                 driver=info.get("description")) for i, (_, info) in enumerate(vulkan_adapters())]
-
-
-def is_available(adapter=None):
-    """True when wgpu is installed and finds a Vulkan GPU (or the adapter selected by `adapter` or by
-    LIGHTPFN_VULKAN_ADAPTER, which may be a CPU driver)."""
-    try:
-        return pick_adapter(adapter) is not None
-    except Exception:
-        return False
-
-
-def __getattr__(name):  # VulkanLightPFN imports torch and the model: load it on first use
-    if name in ("VulkanLightPFN", "VulkanContext"):
-        from lightpfn.vulkan import model
-
-        return getattr(model, name)
-    raise AttributeError(name)
-
-
-__all__ = ["ADAPTER_ENV", "GPU_TYPES", "VulkanLightPFN", "VulkanContext", "adapters", "is_available"]

package/lightpfn/vulkan/engine.py

@@ -1,210 +0,0 @@
-"""Vulkan device, buffers and kernel dispatch through wgpu (WebGPU on the Vulkan backend).
-
-Operations are recorded into one compute pass and submitted on flush() or before a read; WebGPU orders the
-dispatches of a pass and inserts the barriers between them, so a later kernel sees what an earlier one
-wrote. Parameters share one storage buffer, at offsets aligned to the device (at least 256 bytes).
-"""
-
-import math
-import os
-
-import numpy as np
-
-from lightpfn.vulkan import kernels
-
-try:
-    import wgpu
-except ImportError:  # optional dependency: pip install wgpu
-    wgpu = None
-
-ADAPTER_ENV = "LIGHTPFN_VULKAN_ADAPTER"
-GPU_TYPES = ("DiscreteGPU", "IntegratedGPU", "VirtualGPU")
-PARAM_SLOT = 256  # bytes of parameters per operation; u32 63 = slice start
-MAX_GROUPS = 65535
-# GPU work per submission, in floating-point operations: a few ms on a discrete GPU, well under a second on
-# an integrated one. Drivers reset a GPU whose job runs too long (amdgpu's ring timeout, Windows TDR after
-# 2 s), so heavy dispatches are cut into slices of workgroups and submitted in several jobs.
-SUBMIT_FLOPS = 5e10
-SUBMIT_OPS = 4096
-
-
-def vulkan_adapters():
-    """Vulkan adapters seen by wgpu, as (adapter, info) pairs; empty without wgpu or a Vulkan driver."""
-    if wgpu is None:
-        return []
-    try:
-        found = wgpu.gpu.enumerate_adapters_sync()
-    except Exception:  # no loader, no driver
-        return []
-    return [(a, a.info) for a in found if a.info.get("backend_type") == "Vulkan"]
-
-
-def pick_adapter(adapter=None):
-    """adapter: None (LIGHTPFN_VULKAN_ADAPTER, else the first discrete GPU, else any GPU), an index into
-    vulkan_adapters() or a substring of the device name (e.g. "llvmpipe", the CPU driver, for tests)."""
-    found = vulkan_adapters()
-    if adapter is None:
-        adapter = os.environ.get(ADAPTER_ENV) or None
-    if adapter is None:
-        for kind in GPU_TYPES:
-            for a, info in found:
-                if info.get("adapter_type") == kind:
-                    return a
-        return None
-    if isinstance(adapter, int) or str(adapter).isdigit():
-        i = int(adapter)
-        return found[i][0] if 0 <= i < len(found) else None
-    for a, info in found:
-        if str(adapter).lower() in str(info.get("device", "")).lower():
-            return a
-    return None
-
-
-class View:
-    """Rows of a buffer for the kernels: row r = (z, i) with z = r // L, i = r % L, at float offset
-    base + (z // zin) * so + (z % zin) * si + i * ss."""
-
-    __slots__ = ("buf", "base", "L", "zin", "so", "si", "ss")
-
-    def __init__(self, buf, base=0, L=1, zin=1, so=0, si=0, ss=0):
-        self.buf, self.base, self.L, self.zin, self.so, self.si, self.ss = buf, base, L, zin, so, si, ss
-
-    def params(self):
-        return [self.base, self.L, self.zin, self.so, self.si, self.ss]
-
-
-def rows(buf, width, base=0):
-    """Contiguous rows of `width` floats."""
-    return View(buf, base, L=1, zin=1, so=width)
-
-
-class Engine:
-    def __init__(self, adapter=None):
-        if wgpu is None:
-            raise RuntimeError("the Vulkan backend needs wgpu: pip install wgpu")
-        ad = pick_adapter(adapter)
-        if ad is None:
-            raise RuntimeError("no Vulkan adapter found" + (f" matching {adapter!r}" if adapter is not None else ""))
-        self.adapter, self.info = ad, ad.info
-        lim = ad.limits
-        want = ("max-storage-buffer-binding-size", "max-buffer-size", "max-storage-buffers-per-shader-stage",
-                "max-compute-workgroup-storage-size", "max-compute-invocations-per-workgroup",
-                "max-compute-workgroups-per-dimension")
-        self.device = ad.request_device_sync(required_limits={k: lim[k] for k in want if k in lim})
-        lim = self.device.limits
-        self.max_binding = min(lim["max-storage-buffer-binding-size"], lim["max-buffer-size"], (1 << 34) - 16)
-        self.max_groups = min(MAX_GROUPS, lim["max-compute-workgroups-per-dimension"])
-        self.param_slot = max(PARAM_SLOT, lim["min-storage-buffer-offset-alignment"])
-        self.submit_ops = min(SUBMIT_OPS, self.max_binding // self.param_slot)
-        self.usage = wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_SRC | wgpu.BufferUsage.COPY_DST
-        self.pipes = {}
-        self.ops = []  # (pipeline, buffers by binding, params u32, workgroups)
-        self.pending = 0.0  # estimated flops recorded since the last submission
-        self.submit_flops = SUBMIT_FLOPS
-        self.scratch = {}
-        self.sync = self.empty(4)
-
-    # buffers -----------------------------------------------------------------------------------
-    def upload(self, a, dtype=np.float32):
-        a = np.ascontiguousarray(a, dtype=dtype)
-        if a.nbytes == 0:
-            a = np.zeros(4, dtype)
-        if a.nbytes > self.max_binding:
-            raise MemoryError(f"{a.nbytes} bytes exceed the device's {self.max_binding}-byte buffer limit")
-        return self.device.create_buffer_with_data(data=a, usage=self.usage)
-
-    def empty(self, n_floats):
-        size = (max(4, int(n_floats)) + 3) // 4 * 16
-        if size > self.max_binding:
-            raise MemoryError(f"{size} bytes exceed the device's {self.max_binding}-byte buffer limit")
-        return self.device.create_buffer(size=size, usage=self.usage)
-
-    def temp(self, role, n_floats):
-        """A scratch buffer reused by role (dispatches are ordered, so reuse is safe)."""
-        buf = self.scratch.get(role)
-        if buf is None or buf.size < 4 * n_floats:
-            grow = min(buf.size // 2, self.max_binding // 4) if buf is not None else 0  # doubling, within the limit
-            buf = self.scratch[role] = self.empty(max(n_floats, grow))
-        return buf
-
-    def download(self, buf, shape, offset=0):
-        self.flush()
-        count = int(np.prod(shape))
-        if count == 0:
-            return np.empty(shape, np.float32)
-        data = self.device.queue.read_buffer(buf, 4 * offset, 4 * count)
-        return np.frombuffer(data, dtype=np.float32).reshape(shape).copy()
-
-    # kernels -----------------------------------------------------------------------------------
-    def pipeline(self, name, src, flags=(), **consts):
-        key = (name, tuple(sorted(flags)), tuple(sorted(consts.items())))
-        p = self.pipes.get(key)
-        if p is None:
-            code = kernels.render(src, flags, **consts)
-            module = self.device.create_shader_module(code=code)
-            p = self.pipes[key] = self.device.create_compute_pipeline(
-                layout="auto", compute={"module": module, "entry_point": "main"})
-        return p
-
-    def dispatch(self, pipe, buffers, params, groups, flops=0.0):
-        """buffers: {binding: buffer}; params: list of u32 (floats as float32 bits), P[1] and P[63] filled
-        here; groups: workgroups (over two grid dims past 65535); flops: estimated cost, which cuts the
-        dispatch into slices of consecutive workgroups and the work into submissions of ~submit_flops."""
-        if groups <= 0:
-            return
-        params = [int(v) for v in params]
-        if len(params) > 63 or any(v < 0 or v > 0xFFFFFFFF for v in params) or groups > 0xFFFFFFFF:
-            raise ValueError("kernel parameters must fit u32 and leave P[63] for the slice start")
-        parts = max(1, min(groups, math.ceil(flops / self.submit_flops)))
-        step = math.ceil(groups / parts)
-        start = 0
-        while start < groups:
-            count = min(step, groups - start)
-            gx = min(count, self.max_groups)
-            gy = min(count // gx, self.max_groups)
-            if parts == 1 and groups <= self.max_groups ** 2:
-                gy = math.ceil(count / gx)  # the global kernel bound covers a single dispatch's padding
-            else:
-                count = gx * gy  # no padded workgroups may spill into the next slice
-            p = list(params) + [0] * (PARAM_SLOT // 4 - len(params))
-            p[1], p[63] = gx, start
-            cost = flops * count / groups
-            if self.pending + cost > self.submit_flops:
-                self.flush()
-            self.ops.append((pipe, buffers, p, (gx, gy, 1)))
-            self.pending += cost
-            start += count
-            if self.pending >= self.submit_flops or len(self.ops) >= self.submit_ops:
-                self.flush()
-
-    def flush(self):
-        if not self.ops:
-            return
-        slot = self.param_slot // 4
-        P = np.zeros(slot * len(self.ops), np.uint32)
-        for k, (_, _, params, _) in enumerate(self.ops):
-            P[k * slot : k * slot + len(params)] = params
-        pbuf = self.device.create_buffer_with_data(data=P, usage=wgpu.BufferUsage.STORAGE)
-        enc = self.device.create_command_encoder()
-        cp = enc.begin_compute_pass()
-        for k, (pipe, buffers, _, grid) in enumerate(self.ops):
-            entries = [{"binding": 0, "resource": {"buffer": pbuf, "offset": k * self.param_slot, "size": PARAM_SLOT}}]
-            entries += [{"binding": b, "resource": {"buffer": buf, "offset": 0, "size": buf.size}}
-                        for b, buf in sorted(buffers.items())]
-            bg = self.device.create_bind_group(layout=pipe.get_bind_group_layout(0), entries=entries)
-            cp.set_pipeline(pipe)
-            cp.set_bind_group(0, bg)
-            cp.dispatch_workgroups(*grid)
-        cp.end()
-        self.device.queue.submit([enc.finish()])
-        self.ops = []
-        self.pending = 0.0
-
-    def finish(self):
-        """Waits for the submitted work (timing, synchronization with the host)."""
-        self.flush()
-        self.device.queue.read_buffer(self.sync, 0, 4)
-
-
-def f32bits(x):
-    return int(np.array([x], np.float32).view(np.uint32)[0])

package/lightpfn/vulkan/kernels.py

@@ -1,457 +0,0 @@
-"""WGSL compute shaders of the Vulkan backend (compiled to SPIR-V by wgpu when a pipeline is created).
-
-Every kernel reads its parameters from a u32 storage buffer P at offset 0. Tensors are float32 rows
-addressed through views (see engine.View): row r = (z, i) with z = r / L, i = r % L lives at float offset
-base + (z / zin) * so + (z % zin) * si + i * ss. Views let one kernel read a column of cells, a row of
-cells, the first tokens of every row or a broadcast parameter without copies or transposes.
-
-Precision: everything is float32. sin/cos use a Cody-Waite range reduction with minimax polynomials and
-erf the Abramowitz-Stegun 7.1.26 formula (float32 error below 5e-7 on [-32,32]), so results do not depend on the precision
-of the driver's transcendental functions.
-"""
-
-import re
-
-COMMON = """
-@group(0) @binding(0) var<storage, read> P: array<u32>;
-
-struct View { base: u32, L: u32, zin: u32, so: u32, si: u32, ss: u32 }
-
-fn view(o: u32) -> View { return View(P[o], P[o + 1u], P[o + 2u], P[o + 3u], P[o + 4u], P[o + 5u]); }
-fn voff(v: View, z: u32, i: u32) -> u32 { return v.base + (z / v.zin) * v.so + (z % v.zin) * v.si + i * v.ss; }
-fn roff(v: View, r: u32) -> u32 { return voff(v, r / v.L, r % v.L); }
-fn pf(o: u32) -> f32 { return bitcast<f32>(P[o]); }
-
-fn erf_(x: f32) -> f32 {
-  let z = abs(x);
-  let t = 1.0 / (1.0 + 0.3275911 * z);
-  let y = 1.0 - (((((1.061405429 * t - 1.453152027) * t) + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t * exp(-z * z);
-  return select(-y, y, x >= 0.0);
-}
-fn gelu(x: f32) -> f32 { return 0.5 * x * (1.0 + erf_(x * 0.7071067811865476)); }
-fn gelu4(v: vec4<f32>) -> vec4<f32> { return vec4<f32>(gelu(v.x), gelu(v.y), gelu(v.z), gelu(v.w)); }
-fn tanh_(x: f32) -> f32 {
-  let e = exp(2.0 * min(abs(x), 15.0));
-  let t = 1.0 - 2.0 / (e + 1.0);
-  return select(-t, t, x >= 0.0);
-}
-fn sincos(x: f32) -> vec2<f32> {
-  // x = k * pi/2 + r with |r| <= pi/4 (pi/2 split in three parts, Cody-Waite), cephes sinf/cosf polynomials
-  let k = round(x * 0.6366197723675814);
-  var r = x - k * 1.5703125;
-  r = r - k * 4.837512969970703125e-4;
-  r = r - k * 7.54978995489188216e-8;
-  let r2 = r * r;
-  let s = r + r * r2 * (-1.6666654611e-1 + r2 * (8.3321608736e-3 + r2 * -1.9515295891e-4));
-  let c = 1.0 - 0.5 * r2 + r2 * r2 * (4.166664568298827e-2 + r2 * (-1.388731625493765e-3 + r2 * 2.443315711809948e-5));
-  let q = i32(k - 4.0 * floor(k * 0.25));
-  if (q == 0) { return vec2<f32>(s, c); }
-  if (q == 1) { return vec2<f32>(c, -s); }
-  if (q == 2) { return vec2<f32>(-s, -c); }
-  return vec2<f32>(-c, s);
-}
-// P[1]: workgroups along x (2D grids past 65535); P[63]: first workgroup of this slice of the dispatch
-fn wgid(wg: vec3u) -> u32 { return P[63] + wg.x + wg.y * P[1]; }
-"""
-
-# Y[r, :O] = epilogue((X[r, :K] @ W^T)), W (O, K) row-major. 64 x 64 output tiles, 256 threads with 4 x 4
-# outputs each. Options (compile time): NORM scales row r by 1 / rms(X[r]) (a folded RMSNorm), BIAS adds
-# b, GELU applies gelu, RES = "self" adds the old Y, "r" adds rows of a separate view R.
-# P: [0] n_tiles, [1] grid x, [2] N, [3] K, [4] O, [5] eps, [6..11] X view, [12..17] Y view, [18..23] R view
-GEMM = """
-@group(0) @binding(1) var<storage, read> X: array<vec4<f32>>;
-@group(0) @binding(2) var<storage, read> W: array<vec4<f32>>;
-@group(0) @binding(3) var<storage, read_write> Y: array<vec4<f32>>;
-#if BIAS
-@group(0) @binding(4) var<storage, read> Bv: array<vec4<f32>>;
-#endif
-#if RES_R
-@group(0) @binding(5) var<storage, read> R: array<vec4<f32>>;
-#endif
-var<workgroup> xs: array<vec4<f32>, 256>;  // [k 16][row 64 / 4]
-var<workgroup> ws: array<vec4<f32>, 256>;  // [k 16][col 64 / 4]
-var<workgroup> xo: array<u32, 64>;
-var<workgroup> yo: array<u32, 64>;
-var<workgroup> ro: array<u32, 64>;
-
-@compute @workgroup_size(16, 16)
-fn main(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_id) lid: vec3u, @builtin(local_invocation_index) li: u32) {
-  let id = wgid(wg);
-  if (id >= P[0]) { return; }
-  let N = P[2]; let K = P[3]; let O = P[4];
-  let tiles_c = (O + 63u) / 64u;
-  let row0 = (id / tiles_c) * 64u;
-  let col0 = (id % tiles_c) * 64u;
-  if (li < 64u) {
-    let r = min(row0 + li, N - 1u);
-    xo[li] = roff(view(6u), r) / 4u;
-    yo[li] = roff(view(12u), r) / 4u;
-    ro[li] = roff(view(18u), r) / 4u;
-  }
-  workgroupBarrier();
-  var acc: array<vec4<f32>, 4>;
-  var ssq = vec4<f32>(0.0);
-  let lr = li / 4u;  // row (and column) loaded by this thread
-  let lk = li % 4u;  // vec4 of k loaded by this thread
-  let K4 = K / 4u;
-  for (var k4 = 0u; k4 < K4; k4 += 4u) {
-    var xv = vec4<f32>(0.0);
-    if (row0 + lr < N && k4 + lk < K4) { xv = X[xo[lr] + k4 + lk]; }
-    var wv = vec4<f32>(0.0);
-    if (col0 + lr < O && k4 + lk < K4) { wv = W[(col0 + lr) * K4 + k4 + lk]; }
-    for (var c = 0u; c < 4u; c++) {
-      xs[(lk * 4u + c) * 16u + lr / 4u][lr % 4u] = xv[c];
-      ws[(lk * 4u + c) * 16u + lr / 4u][lr % 4u] = wv[c];
-    }
-    workgroupBarrier();
-    for (var kk = 0u; kk < 16u; kk++) {
-      let a = xs[kk * 16u + lid.y];
-      let b = ws[kk * 16u + lid.x];
-      acc[0] += a.x * b;
-      acc[1] += a.y * b;
-      acc[2] += a.z * b;
-      acc[3] += a.w * b;
-#if NORM
-      ssq += a * a;
-#endif
-    }
-    workgroupBarrier();
-  }
-  let c4 = col0 / 4u + lid.x;
-  if (col0 + lid.x * 4u >= O) { return; }
-  for (var i = 0u; i < 4u; i++) {
-    let rl = lid.y * 4u + i;
-    if (row0 + rl >= N) { continue; }
-    var v = acc[i];
-#if NORM
-    v *= inverseSqrt(ssq[i] / f32(K) + pf(5u));
-#endif
-#if BIAS
-    v += Bv[c4];
-#endif
-#if GELU
-    v = gelu4(v);
-#endif
-#if RES_SELF
-    v += Y[yo[rl] + c4];
-#endif
-#if RES_R
-    v += R[ro[rl] + c4];
-#endif
-    Y[yo[rl] + c4] = v;
-  }
-}
-"""
-
-# softmax(q k^T / sqrt(D)) v for every (z, head, query), float32 online softmax (flash attention).
-# Tiled: a workgroup holds 64 queries of one (z, head) and walks the keys in shared-memory tiles of KT.
-# P: [0] n_groups, [1] grid x, [2] Lq, [3] Lk, [4] H, [5] scale, [6..11] q, [12..17] k, [18..23] v,
-#    [24..29] o views, [30] q_h, [31] k_h, [32] v_h, [33] o_h head strides (floats)
-ATTN_TILED = """
-const D4: u32 = {D4}u;
-const KT: u32 = {KT}u;
-@group(0) @binding(1) var<storage, read> Q: array<vec4<f32>>;
-@group(0) @binding(2) var<storage, read> K: array<vec4<f32>>;
-@group(0) @binding(3) var<storage, read> V: array<vec4<f32>>;
-@group(0) @binding(4) var<storage, read_write> O: array<vec4<f32>>;
-var<workgroup> ks: array<vec4<f32>, {KT} * {D4}>;
-var<workgroup> vs: array<vec4<f32>, {KT} * {D4}>;
-
-@compute @workgroup_size(64)
-fn main(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_index) li: u32) {
-  let id = wgid(wg);
-  if (id >= P[0]) { return; }
-  let Lq = P[2]; let Lk = P[3]; let H = P[4];
-  let tiles = (Lq + 63u) / 64u;
-  let zh = id / tiles;
-  let z = zh / H;
-  let h = zh % H;
-  let qi = (id % tiles) * 64u + li;
-  let qv = view(6u); let kv = view(12u); let vv = view(18u); let ov = view(24u);
-  let qb = (voff(qv, z, min(qi, Lq - 1u)) + h * P[30]) / 4u;
-  let scale = pf(5u);
-  var q: array<vec4<f32>, {D4}>;
-  var acc: array<vec4<f32>, {D4}>;
-  for (var d = 0u; d < D4; d++) { q[d] = Q[qb + d] * scale; acc[d] = vec4<f32>(0.0); }
-  var m = -3.0e38;
-  var l = 0.0;
-  for (var t0 = 0u; t0 < Lk; t0 += KT) {
-    for (var e = li; e < KT * D4; e += 64u) {
-      let key = min(t0 + e / D4, Lk - 1u);
-      ks[e] = K[(voff(kv, z, key) + h * P[31]) / 4u + e % D4];
-      vs[e] = V[(voff(vv, z, key) + h * P[32]) / 4u + e % D4];
-    }
-    workgroupBarrier();
-    let nk = min(KT, Lk - t0);
-    var s: array<f32, {KT}>;
-    var mt = m;
-    for (var j = 0u; j < KT; j++) {
-      var dot4 = vec4<f32>(0.0);
-      for (var d = 0u; d < D4; d++) { dot4 += q[d] * ks[j * D4 + d]; }
-      let sj = select(-3.0e38, dot4.x + dot4.y + dot4.z + dot4.w, j < nk);
-      s[j] = sj;
-      mt = max(mt, sj);
-    }
-    let corr = exp(m - mt);
-    l *= corr;
-    for (var d = 0u; d < D4; d++) { acc[d] *= corr; }
-    for (var j = 0u; j < KT; j++) {
-      let p = select(0.0, exp(s[j] - mt), j < nk);
-      l += p;
-      for (var d = 0u; d < D4; d++) { acc[d] += p * vs[j * D4 + d]; }
-    }
-    m = mt;
-    workgroupBarrier();
-  }
-  if (qi < Lq) {
-    let ob = (voff(ov, z, qi) + h * P[33]) / 4u;
-    for (var d = 0u; d < D4; d++) { O[ob + d] = acc[d] / l; }
-  }
-}
-"""
-
-# Same function, one thread per (z, head, query) reading keys from global memory: for few queries per
-# sequence (summary tokens) or few keys. Same parameter layout as ATTN_TILED, [0] = Z * H * Lq.
-ATTN_SMALL = """
-const D4: u32 = {D4}u;
-@group(0) @binding(1) var<storage, read> Q: array<vec4<f32>>;
-@group(0) @binding(2) var<storage, read> K: array<vec4<f32>>;
-@group(0) @binding(3) var<storage, read> V: array<vec4<f32>>;
-@group(0) @binding(4) var<storage, read_write> O: array<vec4<f32>>;
-
-@compute @workgroup_size(64)
-fn main(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_index) li: u32) {
-  let id = wgid(wg) * 64u + li;
-  if (id >= P[0]) { return; }
-  let Lq = P[2]; let Lk = P[3]; let H = P[4];
-  let qi = id % Lq;
-  let zh = id / Lq;
-  let z = zh / H;
-  let h = zh % H;
-  let qv = view(6u); let kv = view(12u); let vv = view(18u); let ov = view(24u);
-  let qb = (voff(qv, z, qi) + h * P[30]) / 4u;
-  let scale = pf(5u);
-  var q: array<vec4<f32>, {D4}>;
-  var acc: array<vec4<f32>, {D4}>;
-  for (var d = 0u; d < D4; d++) { q[d] = Q[qb + d] * scale; acc[d] = vec4<f32>(0.0); }
-  var m = -3.0e38;
-  var l = 0.0;
-  for (var j = 0u; j < Lk; j++) {
-    let kb = (voff(kv, z, j) + h * P[31]) / 4u;
-    var dot4 = vec4<f32>(0.0);
-    for (var d = 0u; d < D4; d++) { dot4 += q[d] * K[kb + d]; }
-    let s = dot4.x + dot4.y + dot4.z + dot4.w;
-    let mt = max(m, s);
-    let corr = exp(m - mt);
-    let p = exp(s - mt);
-    l = l * corr + p;
-    let vb = (voff(vv, z, j) + h * P[32]) / 4u;
-    for (var d = 0u; d < D4; d++) { acc[d] = acc[d] * corr + p * V[vb + d]; }
-    m = mt;
-  }
-  let ob = (voff(ov, z, qi) + h * P[33]) / 4u;
-  for (var d = 0u; d < D4; d++) { O[ob + d] = acc[d] / l; }
-}
-"""
-
-# dst row r = src row r (W4 vec4 per row). P: [0] rows * W4, [1] grid x, [2] W4, [6..11] src, [12..17] dst
-COPY = """
-@group(0) @binding(1) var<storage, read> S: array<vec4<f32>>;
-@group(0) @binding(2) var<storage, read_write> Dst: array<vec4<f32>>;
-
-@compute @workgroup_size(64)
-fn main(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_index) li: u32) {
-  let id = wgid(wg) * 64u + li;
-  if (id >= P[0]) { return; }
-  let r = id / P[2];
-  let c = id % P[2];
-  Dst[roff(view(12u), r) / 4u + c] = S[roff(view(6u), r) / 4u + c];
-}
-"""
-
-# In place on rows of W4 vec4: optional LayerNorm (weight, bias, eps), then optional + emb[slot[z]], where
-# slot index = r / P[4] (the training row of a cell).
-# P: [0] rows, [1] grid x, [2] W4, [3] eps, [4] rows per slot, [6..11] view
-ROWNORM = """
-const W4: u32 = {W4}u;
-@group(0) @binding(1) var<storage, read_write> X: array<vec4<f32>>;
-#if LN
-@group(0) @binding(2) var<storage, read> Lw: array<vec4<f32>>;
-@group(0) @binding(3) var<storage, read> Lb: array<vec4<f32>>;
-#endif
-#if EMB
-@group(0) @binding(4) var<storage, read> Emb: array<vec4<f32>>;
-@group(0) @binding(5) var<storage, read> Slot: array<u32>;
-#endif
-
-@compute @workgroup_size(64)
-fn main(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_index) li: u32) {
-  let r = wgid(wg) * 64u + li;
-  if (r >= P[0]) { return; }
-  let v = view(6u);
-  let b = roff(v, r) / 4u;
-#if LN
-  var x: array<vec4<f32>, {W4}>;
-  var s = vec4<f32>(0.0);
-  for (var c = 0u; c < W4; c++) { x[c] = X[b + c]; s += x[c]; }
-  let mean = (s.x + s.y + s.z + s.w) / f32(W4 * 4u);
-  var q = vec4<f32>(0.0);
-  for (var c = 0u; c < W4; c++) { let d = x[c] - mean; q += d * d; }
-  let rstd = inverseSqrt((q.x + q.y + q.z + q.w) / f32(W4 * 4u) + pf(3u));
-  for (var c = 0u; c < W4; c++) { x[c] = (x[c] - mean) * rstd * Lw[c] + Lb[c]; }
-#if EMB
-  let e = Slot[r / P[4]] * W4;
-  for (var c = 0u; c < W4; c++) { x[c] += Emb[e + c]; }
-#endif
-  for (var c = 0u; c < W4; c++) { X[b + c] = x[c]; }
-#else
-  let e = Slot[r / P[4]] * W4;
-  for (var c = 0u; c < W4; c++) { X[b + c] += Emb[e + c]; }
-#endif
-}
-"""
-
-# Cell features for the embedding GEMM, one thread per cell of a block of the (B, n, m) table: cells c in
-# (b, ii, jj) order over B x nr rows (from row i0) x mc columns (from column j0). Inputs z, r, nan are the
-# normalized values of the whole table; neighbors are the group offsets (j + o) % m.
-# Output row c, KP floats: fourier [sum_g sin(z_g f), sum_g cos(z_g f)] (2F) or rbf kernels (64), then
-# [z_g, nan_g, sin(r_g pi 2^e), cos(r_g pi 2^e)], zero padded.
-# P: [0] cells, [1] grid x, [2] n, [3] m, [4] i0, [5] nr, [6] j0, [7] mc
-FEATURES = """
-const G: u32 = {G}u;
-const NF: u32 = {NF}u;
-const NE: u32 = {NE}u;
-const KP: u32 = {KP}u;
-const OFFS = array<u32, {G}>({OFFS});
-@group(0) @binding(1) var<storage, read> Zt: array<f32>;
-@group(0) @binding(2) var<storage, read> Rt: array<f32>;
-@group(0) @binding(3) var<storage, read> Nt: array<f32>;
-@group(0) @binding(4) var<storage, read> Fq: array<f32>;
-@group(0) @binding(5) var<storage, read_write> Out: array<f32>;
-
-@compute @workgroup_size(64)
-fn main(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_index) li: u32) {
-  let c = wgid(wg) * 64u + li;
-  if (c >= P[0]) { return; }
-  let n = P[2]; let m = P[3]; let nr = P[5]; let mc = P[7];
-  let b = c / (nr * mc);
-  let rem = c % (nr * mc);
-  let i = P[4] + rem / mc;
-  let j = P[6] + rem % mc;
-  let rowbase = (b * n + i) * m;
-  var z: array<f32, {G}>;
-  var r: array<f32, {G}>;
-  var nn: array<f32, {G}>;
-  for (var g = 0u; g < G; g++) {
-    let jg = (j + OFFS[g]) % m;
-    z[g] = Zt[rowbase + jg];
-    r[g] = Rt[rowbase + jg];
-    nn[g] = Nt[rowbase + jg];
-  }
-  let o = c * KP;
-  var k = 0u;
-#if FOURIER
-  for (var f = 0u; f < NF; f++) {
-    var s = 0.0;
-    var co = 0.0;
-    for (var g = 0u; g < G; g++) {
-      let sc = sincos(z[g] * Fq[g * NF + f]);
-      s += sc.x;
-      co += sc.y;
-    }
-    Out[o + f] = s;
-    Out[o + NF + f] = co;
-  }
-  k = 2u * NF;
-#else
-  for (var t = 0u; t < 64u; t++) {
-    var s = 0.0;
-    for (var g = 0u; g < G; g++) { let d = z[g] - Fq[t]; s += exp(-0.5 * d * d); }
-    Out[o + t] = s;
-  }
-  k = 64u;
-#endif
-  for (var g = 0u; g < G; g++) { Out[o + k + g] = z[g]; Out[o + k + G + g] = nn[g]; }
-  k += 2u * G;
-  for (var g = 0u; g < G; g++) {
-    for (var e = 0u; e < NE; e++) {
-      let sc = sincos(r[g] * 3.141592653589793 * f32(1u << e));
-      Out[o + k + g * NE + e] = sc.x;
-      Out[o + k + G * NE + g * NE + e] = sc.y;
-    }
-  }
-  k += 2u * G * NE;
-  for (; k < KP; k++) { Out[o + k] = 0.0; }
-}
-"""
-
-# Rotary embedding in place on adjacent pairs (2p, 2p + 1) of each head (the layout of LightPFN.folded()),
-# for tokens i >= p0 of each sequence at position i - p0; cos/sin from a table T (pos, D / 2, 2).
-# P: [0] rows * H * D4, [1] grid x, [2] H, [3] p0, [4] head stride, [6..11] view
-ROPE = """
-const D4: u32 = {D4}u;
-@group(0) @binding(1) var<storage, read_write> X: array<vec4<f32>>;
-@group(0) @binding(2) var<storage, read> T: array<vec4<f32>>;
-
-@compute @workgroup_size(64)
-fn main(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_index) li: u32) {
-  let id = wgid(wg) * 64u + li;
-  if (id >= P[0]) { return; }
-  let v = view(6u);
-  let per_row = P[2] * D4;
-  let r = id / per_row;
-  let i = r % v.L;
-  if (i < P[3]) { return; }
-  let h = (id % per_row) / D4;
-  let d = id % D4;
-  let a = (roff(v, r) + h * P[4]) / 4u + d;
-  let x = X[a];
-  let cs = T[(i - P[3]) * D4 + d];  // (cos, sin) of pairs 2d and 2d + 1
-  X[a] = vec4<f32>(x.x * cs.x - x.y * cs.y, x.x * cs.y + x.y * cs.x,
-                   x.z * cs.z - x.w * cs.w, x.z * cs.w + x.w * cs.z);
-}
-"""
-
-# Softmax scaling of attention queries in place: q[k] *= base[k % HD] * (1 + tanh(mod[k])) (contiguous).
-# P: [0] total vec4, [1] grid x, [2] HD / 4
-QSCALE = """
-@group(0) @binding(1) var<storage, read_write> Q: array<vec4<f32>>;
-@group(0) @binding(2) var<storage, read> Md: array<vec4<f32>>;
-@group(0) @binding(3) var<storage, read> Bs: array<vec4<f32>>;
-
-@compute @workgroup_size(64)
-fn main(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_index) li: u32) {
-  let id = wgid(wg) * 64u + li;
-  if (id >= P[0]) { return; }
-  let mv = Md[id];
-  let t = vec4<f32>(tanh_(mv.x), tanh_(mv.y), tanh_(mv.z), tanh_(mv.w));
-  Q[id] = Q[id] * Bs[id % P[2]] * (1.0 + t);
-}
-"""
-
-
-def render(src, flags=(), **consts):
-    """Source with `#if NAME` / `#else` / `#endif` blocks resolved (NAME in flags) and {KEY} replaced."""
-    out, stack = [], []
-    for line in (COMMON + src).splitlines():
-        s = line.strip()
-        if s.startswith("#if "):
-            stack.append(s[4:].strip() in flags)
-        elif s == "#else":
-            if not stack:
-                raise ValueError("unmatched #else in shader template")
-            stack[-1] = not stack[-1]
-        elif s == "#endif":
-            if not stack:
-                raise ValueError("unmatched #endif in shader template")
-            stack.pop()
-        elif all(stack):
-            out.append(line)
-    if stack:
-        raise ValueError("unclosed #if in shader template")
-    code = "\n".join(out)
-    for k, v in consts.items():
-        code = code.replace("{" + k + "}", str(v))
-    if re.search(r"\{[A-Z][A-Z0-9_]*\}", code):
-        raise ValueError("missing shader template constant")
-    return code

package/lightpfn/vulkan/model.py

@@ -1,548 +0,0 @@
-"""LightPFN inference on a Vulkan GPU: the computation of LightPFN.folded().encode / predict_logits with
-WGSL kernels. Column statistics and value normalization (sort, searchsorted) stay on the host in torch;
-everything after the normalized values runs on the GPU, and the training-set context (column caches, ICL
-key/value cache, decoder keys) stays in GPU memory between fit and predict.
-"""
-
-import copy
-import math
-from dataclasses import dataclass
-
-import numpy as np
-import torch
-
-from lightpfn.model.layers import Block
-from lightpfn.model.lightpfn import CellEmbedder
-from lightpfn.vulkan import kernels as K
-from lightpfn.vulkan.engine import Engine, View, f32bits, rows
-
-GPU_CHUNK_CELLS = 1 << 20
-
-
-def _eps(norm):
-    return torch.finfo(torch.float32).eps if norm.eps is None else norm.eps
-
-
-def _pad4(n):
-    return (n + 3) // 4 * 4
-
-
-@dataclass
-class VulkanContext:
-    """What prediction needs from the training rows; the caches are GPU buffers."""
-
-    stats: dict
-    col_kv: list
-    col_refine_kv: list | None
-    icl_kv: list
-    dec_k: object
-    onehot: object
-    B: int
-    n: int
-    m: int
-    n_classes: int
-
-
-class _Block:
-    """GPU weights of a folded Block (RMSNorm weights already inside the linear layers)."""
-
-    def __init__(self, eng, blk):
-        a = blk.attn
-        self.H, self.hd = a.n_heads, a.head_dim
-        self.E = a.q.weight.shape[1]
-        self.eps = (_eps(blk.norm_q), _eps(blk.norm_kv), _eps(blk.norm_ff))
-        w = lambda t: eng.upload(t.detach().float().cpu().numpy())  # noqa: E731
-        self.wq, self.wkv, self.wo = w(a.q.weight), w(a.kv.weight), w(a.out.weight)
-        hid = blk.mlp.fc1.weight.shape[0]
-        self.hid = _pad4(hid)  # zero rows/columns: gelu(0) = 0 adds nothing
-        w1 = torch.zeros(self.hid, self.E)
-        w1[:hid] = blk.mlp.fc1.weight
-        w2 = torch.zeros(self.E, self.hid)
-        w2[:, :hid] = blk.mlp.fc2.weight
-        self.w1, self.w2 = w(w1), w(w2)
-        self.scaling = None
-        if a.scaling is not None:
-            self.scaling = _Scaling(eng, a.scaling)
-
-
-class _Scaling:
-    """SoftmaxScaling: base(log n) on the host (a vector per key count), mod(q) on the GPU."""
-
-    def __init__(self, eng, sc):
-        self.base_module = copy.deepcopy(sc.base).float().cpu().eval()
-        self.H, self.hd = sc.n_heads, sc.head_dim
-        w = lambda t: eng.upload(t.detach().float().cpu().numpy())  # noqa: E731
-        lin1, lin2 = sc.mod[0], sc.mod[2]
-        self.hidden = lin1.weight.shape[0]
-        assert self.hidden % 4 == 0
-        self.w1, self.b1, self.w2, self.b2 = w(lin1.weight), w(lin1.bias), w(lin2.weight), w(lin2.bias)
-        self.eng, self.bases = eng, {}
-
-    def base(self, n_keys):
-        buf = self.bases.get(n_keys)
-        if buf is None:
-            with torch.no_grad():
-                b = self.base_module(torch.full((1, 1), math.log(max(n_keys, 2)))).view(-1)
-            buf = self.bases[n_keys] = self.eng.upload(b.numpy())
-        return buf
-
-
-class _ColumnStage:
-    def __init__(self, eng, st):
-        w = lambda t: eng.upload(t.detach().float().cpu().numpy())  # noqa: E731
-        self.y_emb = w(st.y_emb.weight)
-        self.ind = [w(p) for p in st.inducing]
-        self.ind_blocks = [_Block(eng, b) for b in st.ind_blocks]
-        self.cell_blocks = [_Block(eng, b) for b in st.cell_blocks]
-        with torch.no_grad():  # the inducing queries do not depend on the data
-            self.ind_q = [w(b.attn.q(b.norm_q(p))) for b, p in zip(st.ind_blocks, st.inducing)]
-        self.n_ind = st.inducing[0].shape[0]
-
-
-class VulkanLightPFN:
-    """LightPFN on a Vulkan GPU with the encode / predict_logits interface of LightPFN (inputs and logits
-    are CPU torch tensors). model: a LightPFN, folded or not. adapter: see engine.pick_adapter."""
-
-    def __init__(self, model, adapter=None):
-        cfg = model.cfg
-        for dim, heads in ((cfg.col_dim, cfg.col_heads), (cfg.col_dim, cfg.row_heads),
-                           (cfg.icl_dim, cfg.icl_heads), (cfg.icl_dim, cfg.decoder_heads)):
-            if heads <= 0 or dim <= 0 or dim % (4 * heads):
-                raise NotImplementedError("Vulkan attention head dimensions must be positive multiples of four")
-        if not 0 <= cfg.n_ecdf_freq < 32:
-            raise NotImplementedError("Vulkan n_ecdf_freq must be between 0 and 31")
-        if cfg.n_inducing <= 0:
-            raise NotImplementedError("Vulkan inducing token count must be positive")
-        m = model.folded()
-        self.cfg = cfg = m.cfg
-        self.eng = eng = Engine(adapter)
-        self.E, self.D, self.C = cfg.col_dim, cfg.icl_dim, cfg.n_cls
-        if cfg.icl_kv_heads_test not in (1, cfg.icl_heads):
-            raise NotImplementedError("icl_kv_heads_test must be 1 or icl_heads")
-        w = lambda t: eng.upload(t.detach().float().cpu().numpy())  # noqa: E731
-        ce = m.cells
-        self.offsets, self.n_ecdf = tuple(ce.offsets), ce.n_ecdf_freq
-        G = len(self.offsets)
-        meta = G * (2 + 2 * self.n_ecdf)
-        if cfg.cell_embed == "fourier":
-            self.fourier, self.n_freq = True, ce.freq.shape[1]
-            value_w, self.fq = ce.fourier.weight, w(ce.freq)
-        else:
-            self.fourier, self.n_freq = False, 64
-            value_w, self.fq = ce.rbf.weight, w(ce.rbf_centers)
-        k = value_w.shape[1] + meta
-        self.kp = _pad4(k)
-        wc = torch.zeros(self.E, self.kp)
-        wc[:, :k] = torch.cat([value_w, ce.meta.weight], 1)
-        self.w_cells, self.ln_w, self.ln_b, self.ln_eps = w(wc), w(ce.norm.weight), w(ce.norm.bias), ce.norm.eps
-        self.col = _ColumnStage(eng, m.col)
-        self.row_mode, self.rope_base = m.row.mode, m.row.rope_base
-        self.cls = w(m.row.cls)
-        self.row_blocks = [_Block(eng, b) for b in m.row.blocks]
-        self.refine = None
-        if cfg.row_refine:
-            self.refine = dict(summary=w(m.refine.summary), n_summary=m.refine.summary.shape[0],
-                               broadcast=[_Block(eng, b) for b in m.refine.broadcast],
-                               gather=[_Block(eng, b) for b in m.refine.gather])
-            self.col_refine = _ColumnStage(eng, m.col_refine)
-        icl = m.icl
-        self.icl_y_emb, self.thinking = w(icl.y_emb.weight), w(icl.thinking)
-        self.n_thinking = icl.thinking.shape[0]
-        self.icl_blocks = [_Block(eng, b) for b in icl.blocks]
-        self.kv_heads = icl.kv_heads_test
-        dec = m.decoder
-        self.dec_H, self.dec_hd = dec.n_heads, dec.head_dim
-        self.dec_q, self.dec_k = w(dec.q.weight), w(dec.k.weight)
-        self.dec_eps = _eps(icl.norm)
-        self.dec_scaling = _Scaling(eng, dec.scaling)
-        self.rope_tables = {}
-        # Cell limits are upper bounds; train_capacity and the piece costs also bound caches and scratch.
-        self.max_cells = eng.max_binding // (4 * 2 * self.E)
-        self.max_train_cells = eng.max_binding // (4 * self.E)
-
-    def _piece_costs(self, n, m):
-        """Floats per estimator in one column / one row piece."""
-        stages = [self.col] + ([self.col_refine] if self.refine else [])
-        col_width = max(self.kp, 2 * self.E,
-                        *(b.hid for s in stages for b in s.ind_blocks + s.cell_blocks))
-        row_blocks = self.row_blocks + (self.refine["broadcast"] + self.refine["gather"] if self.refine else [])
-        row_width = max(2 * self.E, *(b.hid for b in row_blocks))
-        S = self.refine["n_summary"] if self.refine else 0
-        return max(n, self.col.n_ind) * col_width, max(max(m + self.C, S) * row_width, m * self.kp)
-
-    def _icl_width(self):
-        return max(2 * self.D, self.dec_H * self.dec_scaling.hidden,
-                   *(b.hid for b in self.icl_blocks),
-                   *(b.H * b.scaling.hidden for b in self.icl_blocks))
-
-    def train_capacity(self, n, m):
-        """Estimators whose persistent buffers and smallest pieces fit the binding limit."""
-        col_piece, row_piece = self._piece_costs(n, m)
-        largest = max(n * m * self.E, m * self.col.n_ind * 2 * self.E,
-                      (n + self.n_thinking) * self._icl_width(), col_piece, row_piece)
-        return min(self.max_train_cells // max(1, n * m), self.eng.max_binding // (4 * max(1, largest)))
-
-    # kernels -------------------------------------------------------------------------------------
-    def gemm(self, x, y, w, N, Kd, O, norm_eps=None, bias=None, gelu=False, res=None):
-        """y rows = epilogue(x rows @ w^T) for N rows; res: None, "self" or a View added to the result."""
-        flags = []
-        if norm_eps is not None:
-            flags.append("NORM")
-        if bias is not None:
-            flags.append("BIAS")
-        if gelu:
-            flags.append("GELU")
-        if res == "self":
-            flags.append("RES_SELF")
-        elif res is not None:
-            flags.append("RES_R")
-        assert Kd % 4 == 0 and O % 4 == 0, (Kd, O)
-        tiles = math.ceil(N / 64) * math.ceil(O / 64)
-        rv = res if isinstance(res, View) else y
-        params = [tiles, 0, N, Kd, O, f32bits(norm_eps or 0.0)] + x.params() + y.params() + rv.params()
-        bufs = {1: x.buf, 2: w, 3: y.buf}
-        if bias is not None:
-            bufs[4] = bias
-        if isinstance(res, View):
-            bufs[5] = res.buf
-        self.eng.dispatch(self.eng.pipeline("gemm", K.GEMM, flags), bufs, params, tiles, 2.0 * N * Kd * O)
-
-    def attention(self, q, k, v, o, Z, H, Lq, Lk, hd, q_h, k_h, v_h, o_h):
-        if Z == 0 or Lq == 0:
-            return
-        if Lk == 0:
-            width = H * hd
-            zero = self.eng.upload(np.zeros((Z * Lq, width), np.float32))
-            self.copy(rows(zero, width), o, Z * Lq, width)
-            return
-        small = Lq < 32 or Lk < 16
-        D4 = hd // 4
-        if small:
-            pipe = self.eng.pipeline("attn_small", K.ATTN_SMALL, D4=D4)
-            total = Z * H * Lq
-            groups = math.ceil(total / 64)
-        else:
-            kt = min(64, 16384 // (2 * hd * 4))
-            pipe = self.eng.pipeline("attn_tiled", K.ATTN_TILED, D4=D4, KT=kt)
-            total = math.ceil(Lq / 64) * Z * H
-            groups = total
-        params = [total, 0, Lq, Lk, H, f32bits(hd ** -0.5)] + q.params() + k.params() + v.params() + o.params()
-        params += [q_h, k_h, v_h, o_h]
-        self.eng.dispatch(pipe, {1: q.buf, 2: k.buf, 3: v.buf, 4: o.buf}, params, groups, 4.0 * Z * H * Lq * Lk * hd)
-
-    def copy(self, src, dst, N, width):
-        W4 = width // 4
-        params = [N * W4, 0, W4, 0, 0, 0] + src.params() + dst.params()
-        self.eng.dispatch(self.eng.pipeline("copy", K.COPY), {1: src.buf, 2: dst.buf}, params, math.ceil(N * W4 / 64))
-
-    def rownorm(self, x, N, width, ln=None, emb=None, slots=None, slot_div=1):
-        """In place on N rows: LayerNorm (ln = (w, b, eps)) then + emb[slots[r // slot_div]]."""
-        flags = (["LN"] if ln else []) + (["EMB"] if emb is not None else [])
-        params = [N, 0, width // 4, f32bits(ln[2] if ln else 0.0), slot_div, 0] + x.params()
-        bufs = {1: x.buf}
-        if ln:
-            bufs[2], bufs[3] = ln[0], ln[1]
-        if emb is not None:
-            bufs[4], bufs[5] = emb, slots
-        self.eng.dispatch(self.eng.pipeline("rownorm", K.ROWNORM, flags, W4=width // 4), bufs, params, math.ceil(N / 64))
-
-    def rope(self, x, N, H, hd, head_stride, p0, L):
-        tab = self.rope_table(L, hd)
-        D4 = hd // 4
-        params = [N * H * D4, 0, H, p0, head_stride, 0] + x.params()
-        self.eng.dispatch(self.eng.pipeline("rope", K.ROPE, D4=D4), {1: x.buf, 2: tab}, params, math.ceil(N * H * D4 / 64))
-
-    def rope_table(self, L, hd):
-        """cos/sin of positions 0..L-1 for the adjacent pairs, computed as torch's rope() does."""
-        key = (hd, L)
-        tab = self.rope_tables.get(key)
-        if tab is None:
-            for (d, n), t in self.rope_tables.items():  # a longer table of the same head dim serves too
-                if d == hd and n >= L:
-                    return t
-            inv_freq = self.rope_base ** (-torch.arange(0, hd, 2, dtype=torch.float32) / hd)
-            ang = torch.arange(max(L, 1), dtype=torch.float32)[:, None] * inv_freq[None]
-            tab = self.rope_tables[key] = self.eng.upload(torch.stack([ang.cos(), ang.sin()], -1).numpy())
-        return tab
-
-    def qscale(self, q, N, sc, n_keys):
-        """Queries q (N rows of H * hd, contiguous) times base(log n_keys) * (1 + tanh(mod(q)))."""
-        rows_h = N * sc.H
-        mh = self.eng.temp("mod_h", rows_h * sc.hidden)
-        md = self.eng.temp("mod_o", rows_h * sc.hd)
-        self.gemm(rows(q, sc.hd), rows(mh, sc.hidden), sc.w1, rows_h, sc.hd, sc.hidden, bias=sc.b1, gelu=True)
-        self.gemm(rows(mh, sc.hidden), rows(md, sc.hd), sc.w2, rows_h, sc.hidden, sc.hd, bias=sc.b2)
-        total = rows_h * sc.hd // 4
-        params = [total, 0, sc.H * sc.hd // 4]
-        self.eng.dispatch(self.eng.pipeline("qscale", K.QSCALE), {1: q, 2: md, 3: sc.base(n_keys)}, params,
-                          math.ceil(total / 64))
-
-    def mlp(self, blk, x, N):
-        hid = self.eng.temp("hid", N * blk.hid)
-        self.gemm(x, rows(hid, blk.hid), blk.w1, N, blk.E, blk.hid, norm_eps=blk.eps[2], gelu=True)
-        self.gemm(rows(hid, blk.hid), x, blk.w2, N, blk.hid, blk.E, res="self")
-
-    # stages --------------------------------------------------------------------------------------
-    def embed(self, zrn, x, B, n, m, i0, nr, j0, mc):
-        """Cell embeddings of rows i0..i0+nr, columns j0..j0+mc of the (B, n, m) table into view x."""
-        N = B * nr * mc
-        feat = self.eng.temp("feat", N * self.kp)
-        G = len(self.offsets)
-        consts = dict(G=G, NF=self.n_freq, NE=self.n_ecdf, KP=self.kp,
-                      OFFS=", ".join(f"{o % m}u" for o in self.offsets))
-        pipe = self.eng.pipeline("features", K.FEATURES, ["FOURIER"] if self.fourier else [], **consts)
-        params = [N, 0, n, m, i0, nr, j0, mc]
-        self.eng.dispatch(pipe, {1: zrn[0], 2: zrn[1], 3: zrn[2], 4: self.fq, 5: feat}, params, math.ceil(N / 64))
-        self.gemm(rows(feat, self.kp), x, self.w_cells, N, self.kp, self.E)
-
-    def column_context(self, st, main, B, n, m, cols, slots, zrn=None):
-        """ColumnStage.context on pieces of `cols` columns of the (B, n, m, E) buffer `main` (embedded
-        first when zrn is given); returns the (B * m, n_ind, 2E) key/value caches of the cell blocks."""
-        E, I = self.E, st.n_ind
-        caches = [self.eng.empty(B * m * I * 2 * E) for _ in st.cell_blocks]
-        for j0 in range(0, m, cols):
-            mc = min(cols, m - j0)
-            N = B * n * mc
-            x = View(main, j0 * E, L=mc, zin=1, so=m * E, si=0, ss=E)  # rows (b * n + i, jj)
-            ln = None
-            if zrn is not None:
-                self.embed(zrn, x, B, n, m, 0, n, j0, mc)
-                ln = (self.ln_w, self.ln_b, self.ln_eps)
-            self.rownorm(x, N, E, ln=ln, emb=st.y_emb, slots=slots, slot_div=mc)
-            for ind, ind_q, ib, cb, cache in zip(st.ind, st.ind_q, st.ind_blocks, st.cell_blocks, caches):
-                H, hd = ib.H, ib.hd
-                # inducing points read the cells of their column: Z = B * mc sequences of n keys
-                kv = self.eng.temp("kv", N * 2 * E)
-                self.gemm(x, rows(kv, 2 * E), ib.wkv, N, E, 2 * E, norm_eps=ib.eps[1])
-                kview = lambda base: View(kv, base, L=n, zin=mc, so=n * mc * 2 * E, si=2 * E, ss=mc * 2 * E)  # noqa: E731
-                o = self.eng.temp("o_ind", B * mc * I * E)
-                self.attention(View(ind_q, 0, L=I, zin=1, so=0, si=0, ss=E), kview(0), kview(E),
-                               View(o, 0, L=I, zin=1, so=I * E, ss=E), B * mc, H, I, n, hd, hd, hd, hd, hd)
-                h = self.eng.temp("h_ind", B * mc * I * E)
-                self.gemm(rows(o, E), rows(h, E), ib.wo, B * mc * I, E, E,
-                          res=View(ind, 0, L=I, zin=1, so=0, si=0, ss=E))
-                self.mlp(ib, rows(h, E), B * mc * I)
-                cview = lambda base: View(cache, j0 * I * 2 * E + base, L=I, zin=mc, so=m * I * 2 * E,  # noqa: E731
-                                          si=I * 2 * E, ss=2 * E)
-                self.gemm(rows(h, E), cview(0), cb.wkv, B * mc * I, E, 2 * E, norm_eps=cb.eps[1])
-                self.cell_block(cb, x, N, B * mc, n, mc, n * mc, cview(0), cview(E), I)
-        return caches
-
-    def cell_block(self, cb, x, N, Z, L, zin, so_rows, kview, vview, I):
-        """Cells attend to the inducing states of their column (Block.forward with cached keys). x rows
-        are in (z_outer * L + i, jj) order with zin columns per sequence group."""
-        E = self.E
-        q = self.eng.temp("q", N * E)
-        self.gemm(x, rows(q, E), cb.wq, N, E, E, norm_eps=cb.eps[0])
-        qv = View(q, 0, L=L, zin=zin, so=so_rows * E, si=E, ss=zin * E)
-        o = self.eng.temp("o", N * E)
-        ov = View(o, 0, L=L, zin=zin, so=so_rows * E, si=E, ss=zin * E)
-        self.attention(qv, kview, vview, ov, Z, cb.H, L, I, cb.hd, cb.hd, cb.hd, cb.hd, cb.hd)
-        self.gemm(rows(o, E), x, cb.wo, N, E, E, res="self")
-        self.mlp(cb, x, N)
-
-    def column_query(self, st, caches, x, B, nr, m):
-        """ColumnStage.query on test cells x: rows (b * nr + ii, j), contiguous (B, nr, m, E)."""
-        E, I = self.E, st.n_ind
-        N = B * nr * m
-        for cb, cache in zip(st.cell_blocks, caches):
-            cview = lambda base: View(cache, base, L=I, zin=1, so=I * 2 * E, si=0, ss=2 * E)  # noqa: E731
-            self.cell_block(cb, x, N, B * m, nr, m, nr * m, cview(0), cview(E), I)
-
-    def refine_rows(self, x, R, m):
-        """RowRefinement on R rows of m cells, x a view with L = m (row z, cell i)."""
-        E, rf = self.E, self.refine
-        S = rf["n_summary"]
-        N = R * m
-        s = self.eng.temp("s", R * S * E)
-        self.copy(View(rf["summary"], 0, L=S, zin=1, so=0, si=0, ss=E), rows(s, E), R * S, E)
-        for i, bc in enumerate(rf["broadcast"]):
-            H, hd = bc.H, bc.hd
-            kvs = self.eng.temp("kvs", R * S * 2 * E)
-            self.gemm(rows(s, E), rows(kvs, 2 * E), bc.wkv, R * S, E, 2 * E, norm_eps=bc.eps[1])
-            q = self.eng.temp("q", N * E)
-            self.gemm(x, rows(q, E), bc.wq, N, E, E, norm_eps=bc.eps[0])
-            qv = View(q, 0, L=m, zin=1, so=m * E, ss=E)
-            self.rope(qv, N, H, hd, hd, 0, m)
-            o = self.eng.temp("o", N * E)
-            kv_s = lambda base: View(kvs, base, L=S, zin=1, so=S * 2 * E, ss=2 * E)  # noqa: E731
-            self.attention(qv, kv_s(0), kv_s(E), View(o, 0, L=m, zin=1, so=m * E, ss=E), R, H, m, S, hd, hd, hd, hd, hd)
-            self.gemm(rows(o, E), x, bc.wo, N, E, E, res="self")
-            self.mlp(bc, x, N)
-            if i < len(rf["gather"]):
-                g = rf["gather"][i]
-                kvx = self.eng.temp("kv", N * 2 * E)
-                self.gemm(x, rows(kvx, 2 * E), g.wkv, N, E, 2 * E, norm_eps=g.eps[1])
-                kx = lambda base: View(kvx, base, L=m, zin=1, so=m * 2 * E, ss=2 * E)  # noqa: E731
-                self.rope(kx(0), N, H, hd, hd, 0, m)
-                qs = self.eng.temp("qs", R * S * E)
-                self.gemm(rows(s, E), rows(qs, E), g.wq, R * S, E, E, norm_eps=g.eps[0])
-                os_ = self.eng.temp("os", R * S * E)
-                sv = lambda buf: View(buf, 0, L=S, zin=1, so=S * E, ss=E)  # noqa: E731
-                self.attention(sv(qs), kx(0), kx(E), sv(os_), R, H, S, m, hd, hd, hd, hd, hd)
-                self.gemm(rows(os_, E), rows(s, E), g.wo, R * S, E, E, res="self")
-                self.mlp(g, rows(s, E), R * S)
-
-    def row_stage(self, x, R, m, out):
-        """RowStage on R rows of m cells (x a view with L = m); writes the R rows of C * E floats to view out."""
-        E, C = self.E, self.C
-        T = C + m
-        xr = self.eng.temp("xr", R * T * E)
-        self.copy(View(self.cls, 0, L=C, zin=1, so=0, si=0, ss=E), View(xr, 0, L=C, zin=1, so=T * E, ss=E), R * C, E)
-        self.copy(x, View(xr, C * E, L=m, zin=1, so=T * E, ss=E), R * m, E)
-        tok = lambda buf, w, base=0: View(buf, base, L=T, zin=1, so=T * w, ss=w)  # noqa: E731
-        for blk in self.row_blocks:
-            H, hd = blk.H, blk.hd
-            kv = self.eng.temp("kv", R * T * 2 * E)
-            self.gemm(rows(xr, E), rows(kv, 2 * E), blk.wkv, R * T, E, 2 * E, norm_eps=blk.eps[1])
-            self.rope(tok(kv, 2 * E), R * T, H, hd, hd, C, m)
-            if self.row_mode == "summary":
-                sq = View(xr, 0, L=C, zin=1, so=T * E, ss=E)
-                q = self.eng.temp("q", R * C * E)
-                self.gemm(sq, rows(q, E), blk.wq, R * C, E, E, norm_eps=blk.eps[0])
-                o = self.eng.temp("o", R * C * E)
-                cv = lambda buf: View(buf, 0, L=C, zin=1, so=C * E, ss=E)  # noqa: E731
-                self.attention(cv(q), tok(kv, 2 * E), tok(kv, 2 * E, E), cv(o), R, H, C, T, hd, hd, hd, hd, hd)
-                self.gemm(rows(o, E), sq, blk.wo, R * C, E, E, res="self")
-                self.mlp(blk, sq, R * C)
-            else:
-                q = self.eng.temp("q", R * T * E)
-                self.gemm(rows(xr, E), rows(q, E), blk.wq, R * T, E, E, norm_eps=blk.eps[0])
-                self.rope(tok(q, E), R * T, H, hd, hd, C, m)
-                o = self.eng.temp("o", R * T * E)
-                self.attention(tok(q, E), tok(kv, 2 * E), tok(kv, 2 * E, E), tok(o, E), R, H, T, T, hd, hd, hd, hd, hd)
-                self.gemm(rows(o, E), rows(xr, E), blk.wo, R * T, E, E, res="self")
-                self.mlp(blk, rows(xr, E), R * T)
-        self.copy(View(xr, 0, L=1, zin=1, so=T * E), out, R, C * E)
-
-    def icl_block(self, blk, x, B, Lq, kview, vview, Lk, k_h, v_h, kv_heads_cache=None):
-        D = self.D
-        N = B * Lq
-        q = self.eng.temp("q", N * D)
-        self.gemm(rows(x, D), rows(q, D), blk.wq, N, D, D, norm_eps=blk.eps[0])
-        self.qscale(q, N, blk.scaling, Lk)
-        o = self.eng.temp("o", N * D)
-        qv = lambda buf: View(buf, 0, L=Lq, zin=1, so=Lq * D, ss=D)  # noqa: E731
-        self.attention(qv(q), kview, vview, qv(o), B, blk.H, Lq, Lk, blk.hd, blk.hd, k_h, v_h, blk.hd)
-        self.gemm(rows(o, D), rows(x, D), blk.wo, N, D, D, res="self")
-        self.mlp(blk, rows(x, D), N)
-
-    # interface -----------------------------------------------------------------------------------
-    def _upload_normalized(self, X, stats):
-        z, r, nan = CellEmbedder.normalize(X, stats)
-        return tuple(self.eng.upload(t.contiguous().numpy()) for t in (z, r, nan))
-
-    def _chunk(self, chunk_cells):
-        return max(1, min(chunk_cells or GPU_CHUNK_CELLS, self.max_cells))
-
-    @torch.inference_mode()
-    def encode(self, X_train, y_train, d=None, slots=None, n_classes=None, has_padding=None, chunk_cells=None):
-        if d is not None:
-            raise NotImplementedError("the Vulkan backend does not take zero-padded feature counts (d)")
-        X = torch.as_tensor(X_train).float().cpu()
-        y = torch.as_tensor(y_train).long().cpu()
-        B, n, m = X.shape
-        if B == 0 or m == 0:
-            raise ValueError("Vulkan encode needs a nonempty batch and at least one feature")
-        if n == 0 and n_classes is None:
-            raise ValueError("n_classes is required for an empty training context")
-        n_classes = int(y.max()) + 1 if n_classes is None else n_classes
-        if n_classes > self.cfg.label_slots:
-            raise ValueError(f"{n_classes} classes, the model supports {self.cfg.label_slots}")
-        slots = torch.arange(self.cfg.label_slots).expand(B, -1) if slots is None else torch.as_tensor(slots).long().cpu()
-        y_slots = torch.gather(slots, 1, y)
-        if B > self.train_capacity(n, m):
-            raise MemoryError(f"{B} x {n} x {m} training cells exceed a Vulkan cache/scratch buffer limit")
-        eng, E, D = self.eng, self.E, self.D
-        stats = CellEmbedder.stats(X)
-        zrn = self._upload_normalized(X, stats)
-        slot_buf = eng.upload(y_slots.numpy().reshape(-1), np.uint32)
-        chunk = self._chunk(chunk_cells)
-        col_cost, row_cost = self._piece_costs(n, m)
-        cols = min(max(1, chunk // (B * max(n, 1))), eng.max_binding // (4 * B * col_cost))
-        nrows = min(max(1, chunk // (B * m)), eng.max_binding // (4 * B * row_cost))
-        main = eng.empty(B * n * m * E)
-        col_kv = self.column_context(self.col, main, B, n, m, cols, slot_buf, zrn)
-        row_view = lambda i0, nr: View(main, i0 * m * E, L=m, zin=nr, so=n * m * E, si=m * E, ss=E)  # noqa: E731
-        col_refine_kv = None
-        if self.refine is not None:
-            for i0 in range(0, n, nrows):
-                nr = min(nrows, n - i0)
-                self.refine_rows(row_view(i0, nr), B * nr, m)
-            col_refine_kv = self.column_context(self.col_refine, main, B, n, m, cols, slot_buf)
-        rows_buf = eng.empty(B * n * D)
-        for i0 in range(0, n, nrows):
-            nr = min(nrows, n - i0)
-            self.row_stage(row_view(i0, nr), B * nr, m, View(rows_buf, i0 * D, L=1, zin=nr, so=n * D, si=D))
-        del main
-        # ICL over thinking rows + training rows
-        T = self.n_thinking
-        L = T + n
-        x = eng.empty(B * L * D)
-        self.copy(View(self.thinking, 0, L=T, zin=1, so=0, si=0, ss=D), View(x, 0, L=T, zin=1, so=L * D, ss=D), B * T, D)
-        train = View(x, T * D, L=n, zin=1, so=L * D, ss=D)
-        self.copy(rows(rows_buf, D), train, B * n, D)
-        self.rownorm(train, B * n, D, emb=self.icl_y_emb, slots=slot_buf, slot_div=1)
-        del rows_buf
-        kvh = self.kv_heads * self.icl_blocks[0].hd
-        icl_kv = []
-        for blk in self.icl_blocks:
-            kv = eng.temp("kv_icl", B * L * 2 * D)
-            self.gemm(rows(x, D), rows(kv, 2 * D), blk.wkv, B * L, D, 2 * D, norm_eps=blk.eps[1])
-            cache = eng.empty(B * L * 2 * kvh)
-            self.copy(rows(kv, 2 * D), rows(cache, 2 * kvh), B * L, kvh)
-            self.copy(rows(kv, 2 * D, D), rows(cache, 2 * kvh, kvh), B * L, kvh)
-            kview = lambda base: View(kv, base, L=L, zin=1, so=L * 2 * D, ss=2 * D)  # noqa: E731
-            self.icl_block(blk, x, B, L, kview(0), kview(D), L, blk.hd, blk.hd)
-            icl_kv.append(cache)
-        dec_k = eng.empty(B * n * D)
-        self.gemm(train, rows(dec_k, D), self.dec_k, B * n, D, D, norm_eps=self.dec_eps)
-        onehot = torch.nn.functional.one_hot(y, self.dec_hd).float()
-        ctx = VulkanContext(stats, col_kv, col_refine_kv, icl_kv, dec_k, eng.upload(onehot.numpy()), B, n, m, n_classes)
-        eng.flush()
-        return ctx
-
-    @torch.inference_mode()
-    def predict_logits(self, ctx, X_test, chunk_cells=None):
-        X = torch.as_tensor(X_test).float().cpu()
-        B, M, m = X.shape
-        if (B, m) != (ctx.B, ctx.m):
-            raise ValueError(f"test batch {(B, m)} does not match the context {(ctx.B, ctx.m)}")
-        if M == 0:
-            return torch.zeros(B, 0, ctx.n_classes)
-        eng, E, D = self.eng, self.E, self.D
-        col_cost, row_cost = self._piece_costs(1, m)
-        row_cost = max(row_cost, m * col_cost // self.col.n_ind)
-        if 4 * B * max(M * m, M * self._icl_width(), row_cost) > eng.max_binding:
-            raise MemoryError("query cache/scratch buffer exceeds the Vulkan limit; use fewer query rows")
-        zrn = self._upload_normalized(X, ctx.stats)
-        step = max(1, self._chunk(chunk_cells) // (B * m))
-        step = min(step, eng.max_binding // (4 * B * row_cost))
-        rows_buf = eng.empty(B * M * D)
-        for i0 in range(0, M, step):
-            nr = min(step, M - i0)
-            cells = eng.temp("test_cells", B * nr * m * E)
-            x = View(cells, 0, L=m, zin=1, so=m * E, ss=E)
-            self.embed(zrn, x, B, M, m, i0, nr, 0, m)
-            self.rownorm(x, B * nr * m, E, ln=(self.ln_w, self.ln_b, self.ln_eps))
-            self.column_query(self.col, ctx.col_kv, x, B, nr, m)
-            if self.refine is not None:
-                self.refine_rows(x, B * nr, m)
-                self.column_query(self.col_refine, ctx.col_refine_kv, x, B, nr, m)
-            self.row_stage(x, B * nr, m, View(rows_buf, i0 * D, L=1, zin=nr, so=M * D, si=D))
-        Lc = self.n_thinking + ctx.n
-        kvh = self.kv_heads * self.icl_blocks[0].hd
-        for blk, cache in zip(self.icl_blocks, ctx.icl_kv):
-            kview = lambda base: View(cache, base, L=Lc, zin=1, so=Lc * 2 * kvh, ss=2 * kvh)  # noqa: E731
-            head = 0 if self.kv_heads == 1 else blk.hd
-            self.icl_block(blk, rows_buf, B, M, kview(0), kview(kvh), Lc, head, head)
-        H, hd = self.dec_H, self.dec_hd
-        q = eng.temp("q", B * M * D)
-        self.gemm(rows(rows_buf, D), rows(q, D), self.dec_q, B * M, D, D, norm_eps=self.dec_eps)
-        self.qscale(q, B * M, self.dec_scaling, ctx.n)
-        o = eng.temp("o", B * M * D)
-        qv = lambda buf: View(buf, 0, L=M, zin=1, so=M * D, ss=D)  # noqa: E731
-        self.attention(qv(q), View(ctx.dec_k, 0, L=ctx.n, zin=1, so=ctx.n * D, ss=D),
-                       View(ctx.onehot, 0, L=ctx.n, zin=1, so=ctx.n * hd, ss=hd), qv(o), B, H, M, ctx.n, hd, hd, hd, 0, hd)
-        p = torch.from_numpy(eng.download(o, (B, M, H, hd))).mean(2)[..., : ctx.n_classes]
-        return torch.log(p.clamp(min=1e-5) + 3e-5)

package/pyproject.toml

@@ -1,43 +0,0 @@
-[build-system]
-requires = ["hatchling>=1.27"]
-build-backend = "hatchling.build"
-
-[project]
-name = "lightpfn"
-version = "0.1.0"
-description = "A compact tabular classifier pretrained only on synthetic data"
-readme = "docs/PACKAGE_README.md"
-requires-python = ">=3.10"
-license = "Apache-2.0"
-license-files = ["LICENSE", "NOTICE"]
-authors = [{name = "Giorgio", email = "247403232+GioOtto@users.noreply.github.com"}]
-dependencies = ["numpy>=1.24", "torch>=2.6"]
-classifiers = [
-  "Development Status :: 3 - Alpha",
-  "Intended Audience :: Science/Research",
-  "Programming Language :: Python :: 3",
-  "Operating System :: OS Independent",
-  "Topic :: Scientific/Engineering :: Artificial Intelligence",
-]
-
-[project.optional-dependencies]
-sklearn = ["scikit-learn>=1.6"]
-hf = ["huggingface-hub>=0.27", "safetensors>=0.5"]
-vulkan = ["wgpu>=0.32,<0.33"]
-dev = ["pytest>=8", "build>=1.2", "hatchling>=1.27"]
-
-[project.urls]
-Repository = "https://github.com/GioOtto/gioPFN"
-Models = "https://huggingface.co/ueuegio/LightPFN"
-
-[tool.hatch.build.targets.wheel]
-packages = ["lightpfn"]
-exclude = ["lightpfn/train.py", "lightpfn/prior", "lightpfn/eval", "lightpfn/model/ccmm.py"]
-
-[tool.hatch.build.targets.sdist]
-include = [
-  "/lightpfn", "/pyproject.toml", "/LICENSE", "/NOTICE",
-  "/docs/PACKAGE_README.md", "/tests/test_release.py", "/tests/test_sklearn_api.py",
-  "/docs/DEPENDENCIES.md",
-]
-exclude = ["lightpfn/train.py", "lightpfn/prior", "lightpfn/eval", "lightpfn/model/ccmm.py", "**/__pycache__"]

package/tests/test_release.py

@@ -1,94 +0,0 @@
-"""Release checkpoints must be safe, exact, and usable without research dependencies."""
-
-import json
-import pickle
-import subprocess
-import sys
-from dataclasses import asdict
-
-import numpy as np
-import pytest
-import torch
-
-from lightpfn import Config, LightPFN, load_model, save_model
-from lightpfn.checkpoint import load_pretrained
-
-
-def small_model():
-    torch.manual_seed(13)
-    model = LightPFN(Config(col_dim=16, col_blocks=1, col_heads=4, n_inducing=4,
-                            row_blocks=1, row_heads=4, n_cls=2, icl_blocks=2,
-                            icl_heads=4, decoder_heads=2, n_thinking=2, n_freq=4, n_ecdf_freq=2)).eval()
-    for parameter in model.parameters():
-        torch.nn.init.normal_(parameter, std=0.05)
-    return model
-
-
-def test_legacy_and_safetensors_roundtrip_predictions(tmp_path):
-    model = small_model()
-    path = tmp_path / "ema.pt"
-    torch.save(dict(config=asdict(model.cfg), ema=model.state_dict()), path)
-    legacy = load_model(path)
-    release = save_model(legacy, tmp_path / "release")
-    restored = load_model(release)
-    assert restored.cfg == model.cfg
-    assert set(json.loads((release / "config.json").read_text())) == set(asdict(model.cfg))
-    g = torch.Generator().manual_seed(3)
-    X = torch.randn(1, 22, 3, generator=g)
-    X[0, 1, 0] = float("nan")
-    y = torch.arange(16).remainder(3).unsqueeze(0)
-    with torch.inference_mode():
-        expected = model(X, y, n_classes=3)
-        torch.testing.assert_close(legacy(X, y, n_classes=3), expected, rtol=0, atol=0)
-        torch.testing.assert_close(restored(X, y, n_classes=3), expected, rtol=0, atol=0)
-
-
-def test_pickle_payload_is_rejected_without_execution(tmp_path):
-    class Payload:
-        def __reduce__(self):
-            return eval, ("__import__('pathlib').Path(%r).touch()" % str(tmp_path / "executed"),)
-    path = tmp_path / "untrusted.pt"
-    torch.save(dict(config={}, model=Payload()), path)
-    with pytest.raises(pickle.UnpicklingError):
-        load_model(path)
-    assert not (tmp_path / "executed").exists()
-
-
-def test_reject_malformed_and_folded_checkpoints(tmp_path):
-    path = tmp_path / "bad.pt"
-    torch.save(dict(config={}, model={"not_a_tensor": 5}), path)
-    with pytest.raises(ValueError, match="tensor state"):
-        load_model(path)
-    with pytest.raises(ValueError, match="folded"):
-        save_model(small_model().folded(), tmp_path)
-
-
-def test_hub_files_use_same_resolved_commit(tmp_path, monkeypatch):
-    revision = "a" * 40
-    folder = save_model(small_model(), tmp_path / "snapshots" / revision)
-    calls = []
-    def download(filename, **kwargs):
-        calls.append((filename, kwargs))
-        return str(folder / filename)
-    monkeypatch.setattr("huggingface_hub.hf_hub_download", download)
-    load_pretrained(repo_id="test/model", revision="main", local_files_only=True)
-    assert calls[0][1]["revision"] == "main"
-    assert calls[1][1]["revision"] == revision
-    assert all(options["local_files_only"] for _, options in calls)
-    with pytest.raises(ValueError, match="revision"):
-        load_pretrained(repo_id="test/other")
-
-
-def test_core_import_does_not_load_optional_dependencies():
-    code = """
-import sys
-from importlib.abc import MetaPathFinder
-class Block(MetaPathFinder):
-    def find_spec(self, fullname, path=None, target=None):
-        if fullname.split('.')[0] in {'sklearn', 'tabicl', 'pandas', 'scipy', 'wgpu', 'huggingface_hub', 'safetensors'}:
-            raise ModuleNotFoundError(fullname, name=fullname)
-sys.meta_path.insert(0, Block())
-from lightpfn import LightPFN, Config, load_model
-assert LightPFN(Config()).cfg.label_slots == 16
-"""
-    subprocess.run([sys.executable, "-c", code], check=True, capture_output=True, text=True)

package/tests/test_sklearn_api.py

@@ -1,124 +0,0 @@
-"""Sklearn integration exercises real inference, cloning and input validation."""
-
-import numpy as np
-from dataclasses import asdict
-import pytest
-import torch
-from sklearn.base import clone, is_classifier
-from sklearn.exceptions import NotFittedError
-from sklearn.model_selection import GridSearchCV
-from sklearn.pipeline import Pipeline
-from sklearn.preprocessing import StandardScaler
-from sklearn.utils.estimator_checks import check_estimator
-
-from lightpfn import LightPFNClassifier
-from test_release import small_model
-
-torch.set_num_threads(4)
-
-
-@pytest.fixture
-def task():
-    rng = np.random.default_rng(8)
-    X = rng.normal(size=(32, 3)).astype(np.float32)
-    y = np.where(X[:, 0] > 0, "yes", "no")
-    return X, y
-
-
-def test_construction_and_clone_do_not_load_weights(monkeypatch):
-    def fail(*args, **kwargs):
-        raise AssertionError("constructor performed IO")
-    monkeypatch.setattr("lightpfn.sklearn.load_pretrained", fail)
-    monkeypatch.setattr("lightpfn.sklearn.resolve_device", fail)
-    clf = LightPFNClassifier(device="auto", n_estimators=3, random_state=12)
-    cloned = clone(clf)
-    assert cloned.get_params() == clf.get_params()
-    assert is_classifier(cloned)
-    assert not hasattr(cloned, "net_")
-    with pytest.raises(NotFittedError):
-        cloned.predict([[0, 1]])
-
-
-def test_pipeline_grid_search_and_set_params(task):
-    X, y = task
-    clf = LightPFNClassifier(model=small_model(), device="cpu", random_state=3)
-    pipe = Pipeline([("scale", StandardScaler()), ("model", clf)])
-    search = GridSearchCV(pipe, {"model__n_estimators": [1, 2]}, cv=2).fit(X, y)
-    assert search.predict(X[:4]).shape == (4,)
-    fitted = search.best_estimator_.named_steps["model"]
-    assert fitted.n_features_in_ == 3
-    assert 0 <= fitted.score(X, y) <= 1
-    changed = clone(clf).set_params(n_estimators=2).fit(X, y)
-    assert len(changed.members_) == 2
-    assert not hasattr(clone(changed), "members_")
-
-
-def test_input_errors_and_nan_support(task):
-    X, y = task
-    clf = LightPFNClassifier(model=small_model(), device="cpu")
-    with pytest.raises(ValueError, match="inconsistent"):
-        clf.fit(X, y[:-1])
-    with pytest.raises(ValueError, match="Unknown label|continuous"):
-        clf.fit(X, np.linspace(0, 1, len(y)))
-    X[0, 1] = np.nan
-    clf.fit(X, y)
-    P = clf.predict_proba(X)
-    assert np.isfinite(P).all()
-    np.testing.assert_allclose(P.sum(1), 1, atol=1e-6)
-    with pytest.raises(ValueError, match="features"):
-        clf.predict(X[:, :2])
-    with pytest.raises(ValueError, match="infinity"):
-        clf.predict([[0, np.inf, 1]])
-    with pytest.raises(ValueError, match="10 classes"):
-        clf.fit(X, np.arange(len(y)))
-    with pytest.raises(ValueError, match="max_context"):
-        clf.set_params(max_context=1).fit(X, y)
-    with pytest.raises(NotFittedError):
-        clf.predict(X)
-
-
-def test_categorical_mask_warns_and_preserves_numeric_behavior(task):
-    X, y = task
-    model = small_model()
-    plain = LightPFNClassifier(model=model, device="cpu").fit(X, y)
-    other = LightPFNClassifier(model=model, device="cpu")
-    with pytest.warns(FutureWarning, match="native categorical"):
-        other.fit(X, y, cat=np.array([True, False, False]))
-    np.testing.assert_array_equal(plain.predict_proba(X), other.predict_proba(X))
-    with pytest.raises(ValueError, match="boolean mask"):
-        other.fit(X, y, cat=[0])
-
-
-def test_refit_and_single_class(task):
-    X, y = task
-    model = small_model().train()
-    clf = LightPFNClassifier(model=model, device="cpu").fit(X, y)
-    assert model.training and not clf.model_.training
-    clf.fit(X[:, :2], np.full(len(y), "only"))
-    assert clf.n_features_in_ == 2
-    np.testing.assert_array_equal(clf.predict_proba(X[:, :2]), np.ones((len(y), 1)))
-
-
-def test_feature_name_validation(task):
-    pd = pytest.importorskip("pandas")
-    X, y = task
-    frame = pd.DataFrame(X, columns=["a", "b", "c"])
-    clf = LightPFNClassifier(model=small_model(), device="cpu").fit(frame, y)
-    np.testing.assert_array_equal(clf.feature_names_in_, frame.columns)
-    with pytest.raises(ValueError, match="Feature names|feature names"):
-        clf.predict(frame[["c", "b", "a"]])
-
-
-def test_sklearn_estimator_checks(tmp_path):
-    # joblib hashes torch storage identity, so raw nn.Module parameters cannot
-    # satisfy its deepcopy/hash comparison. Exercise the distribution's path API.
-    model = small_model()
-    path = tmp_path / "model.pt"
-    torch.save(dict(config=asdict(model.cfg), model=model.state_dict()), path)
-    check_estimator(
-        LightPFNClassifier(checkpoint=path, device="cpu", n_threads=4),
-        expected_failed_checks={
-            "check_methods_sample_order_invariance": "FP32 GPU/CPU kernels can differ by ~6e-8 after row permutation.",
-            "check_methods_subset_invariance": "FP32 kernels can differ slightly when query batch shapes change.",
-        },
-    )

provenance.json

@@ -1,18 +1,20 @@
-{
-  "package_version": "0.1.0",
-  "source_commit": "ef5d60de31519d54e361ffb8554c6cfc9bfa1835",
-  "source_checkpoint": "final_long/ema_step039250.pt",
-  "source_checkpoint_sha256": "d2142d467a0bf6935f1a620fc216cb8473ba62815a4e52c0dad6e479a43a2c28",
-  "parameters": 4603088,
-  "torch_version": "2.14.1+cpu",
-  "numpy_version": "2.5.3",
-  "license": "Apache-2.0",
-  "weights_format": "safetensors",
-  "dtype": "float32",
-  "hashes": {
-    "config.json": "cc68fb632a4c4406fd8e1a4d18ddf144c94879e5ab4a61e8f191c2ff4014e9d3",
-    "model.safetensors": "a492572bd1892f8af9a92798fa1316de410c49171eae0ad4821caddbb0be2668",
-    "LightPFN_report.pdf": "8bea6f04f76fa4d082d8884445169b20c9e55b901543c205de46d9fa7345d7b0"
-  },
-  "roundtrip": "bitwise identical tensor weights and CPU logits"
-}
+{
+  "package_version": "1.0.0",
+  "source_checkpoint": "final_long/ema_step039250.pt",
+  "source_checkpoint_sha256": "d2142d467a0bf6935f1a620fc216cb8473ba62815a4e52c0dad6e479a43a2c28",
+  "parameters": 4603088,
+  "torch_version": "2.14.1+cpu",
+  "numpy_version": "2.5.3",
+  "license": "Apache-2.0",
+  "weights_format": "safetensors",
+  "dtype": "float32",
+  "hashes": {
+    "config.json": "cc68fb632a4c4406fd8e1a4d18ddf144c94879e5ab4a61e8f191c2ff4014e9d3",
+    "model.safetensors": "a492572bd1892f8af9a92798fa1316de410c49171eae0ad4821caddbb0be2668",
+    "LightPFN_report.pdf": "274be6425d074c984ce49cee56f40a0801421017a2851febe916ac7bbc4718f4"
+  },
+  "roundtrip": "bitwise identical tensor weights and CPU logits",
+  "weights_revision": "bd389ab59a89dd0e05c9ecb7c642c08ee52e9637",
+  "source_repository": "https://github.com/GioOtto/LightPFN",
+  "package": "https://pypi.org/project/LightPFN/"
+}