Fix open bugs: MPS float64 crash (#68), OOF calibration with max_num_rows (#55), sklearn transform_output leak (#58), safetensors dep (#56)

[p4vlos]

This PR fixes four open issues (rebased on current main; an earlier revision also addressed #43, which 1.0.1 has since fixed upstream). Each fix is small and independent; happy to split into separate PRs if preferred.

#68 — PyTorch backend crash on Apple MPS with float64 targets

_predict_step_pytorch (and the context-cache prefill added in #62) moved y to the device before the float64→float32 guard, but MPS rejects float64 tensors at transfer time, so the guard was unreachable on Apple Silicon. The cast now happens host-side before the move in both places (numerically identical on CPU/CUDA). Added a CPU test asserting the model receives float32 targets for float64 (numpy-default) input.

#55 — calibration fit on partially-zeroed OOF probabilities with max_num_rows

With max_num_rows, each ensemble member draws its own row subsample, so on the single-validation-split path the members' OOF rows do not coincide. fit() sliced all members by member 0's validation indices, so _fit_calibration received rows that were all-zero for every other member (rows summing to 1/E, 2/E, …).

predict_oof_proba now records which (member, row) pairs were actually predicted (oof_pred_mask_), and fit() averages each row over the members that predicted it, dropping rows no member predicted. For all previously-correct configurations (no max_num_rows, or full CV without subsampling) the computation is unchanged — the mask is full there, so the count-normalized sum equals the old mean, and the NNLS path (which already forbids max_num_rows) sees identical inputs. Added a test that spies on _fit_calibration and asserts every calibration row is a valid probability vector under max_num_rows + single-split.

#58 — 'DataFrame' object has no attribute 'flatten' under global pandas output

A user-level sklearn.set_config(transform_output="pandas") made the internal transformers (y_scaler_, preprocessors, ColumnTransformer) return DataFrames, crashing the numpy-based pipeline. Public fit/predict/predict_proba/predict_oof* entry points now run under sklearn.config_context(transform_output="default") via a small decorator. (Instance-level set_output pinning was not an option: ColumnTransformer.set_output requires every sub-transformer to support set_output, which the custom encoders here do not.) Added a test running both estimators under the global pandas config.

#56 — NameError: name 'safetensors' is not defined

The published HF checkpoint is safetensors-only (see #53/#77), but the pytorch extra does not declare safetensors, so huggingface_hub's mixin fails with NameError when the package is missing. Added safetensors to the pytorch extra and the pinned requirements.txt.

Testing

• pytest (torch env, no JAX): 26 passed, 30 skipped (JAX-only) — includes 3 new tests.
• Negative control: with the source fix reverted, the new max_num_rows + calibration: single-split path fits calibration on partially-zeroed OOF probabilities #55 and Regression example not working AttributeError: 'DataFrame' object has no attribute 'flatten' #58 tests fail (the MPS-cast test is contract-only, since the crash needs MPS hardware).

🤖 Generated with Claude Code

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.