Make Poisson sampling opt-in in the Keras DP path

[staticpayload]

What this PR changes

This PR updates the Keras DP path so Poisson sampling can be done directly inside fit() without breaking the existing API.

The main change is an opt-in flag, DPKerasConfig.poisson_sampling_in_fit, which defaults to False. When it is enabled, the wrapper builds a Poisson-sampled keras.utils.PyDataset from random-access per-example arrays, validates the input sizes against the DP config, and rejects validation_split because the accountant needs the exact training-set size.

Main changes

• Add poisson_sampling_in_fit to DPKerasConfig, defaulting to False for backward compatibility.
• Build Poisson-sampled batches internally only when that flag is enabled.
• Validate random-access array inputs and ensure x, y, and sample_weight agree on batch size in the Poisson-in-fit path.
• Represent padded examples through zero sample weights and derive the padding mask from sample_weight inside the DP training path.
• Update the Keras example and docs to show the opt-in Poisson path explicitly.
• Expand Keras tests around batch construction, padding, validation errors, and fit-argument rewriting.
• Add the empty-input padding edge case coverage in batch_selection.

Why this shape

The goal here is to make the privacy-sensitive Poisson-sampling path available directly in the Keras wrapper while keeping the older workflows intact. Users who already batch or sample outside the wrapper can keep doing that. Users who want the wrapper to own batch formation can enable the flag and get a path whose sampling behavior matches the DP accounting assumptions.

Verification

I ran the following locally in Python 3.11:

• KERAS_BACKEND=jax python -m pytest tests/keras_api_test.py -q
• python -m pytest tests/batch_selection_test.py -q -k 'pad_to_multiple_of'
• python -m pyink --check --diff jax_privacy/keras_api.py tests/keras_api_test.py
• python -m pylint --rcfile=.pylintrc jax_privacy/keras_api.py

[staticpayload]

@ryan112358

[staticpayload]

Merge?

[ryan112358]

Merge?

We are having one more internal reviewer take a look at this PR before merging. Given the size of the PR, it may take a couple of days.

[dvadym]

Thanks it looks cool. I've left some comments.

[codecov-commenter]

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

❌ Patch coverage is 90.85366% with 15 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (main@aeae005). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
jax_privacy/keras_api.py	91.35%	14 Missing ⚠️
jax_privacy/batch_selection.py	50.00%	1 Missing ⚠️
❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #196   +/-   ##
=======================================
  Coverage        ?   73.60%           
=======================================
  Files           ?       25           
  Lines           ?     3020           
  Branches        ?        0           
=======================================
  Hits            ?     2223           
  Misses          ?      797           
  Partials        ?        0           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[ryan112358]

Please resolve merge conflicts

[dvadym]

Thanks!

[RamSaw]

Hi! With syncing to head you removed all changes that were previously added. Please sync to head carefully without reverting changes.

[staticpayload]

I checked this carefully after the rebase rewrite.

The branch update was only a history rewrite onto the latest upstream/main to clear the merge conflicts. The content of the 8 files that make up this PR is unchanged relative to the pre-rewrite branch head. I compared the rewritten branch against the backed-up pre-rewrite head file-by-file for:

• README.md
• docs/keras_api.rst
• docs/overview.md
• examples/keras_api_example.py
• jax_privacy/batch_selection.py
• jax_privacy/keras_api.py
• tests/batch_selection_test.py
• tests/keras_api_test.py

That comparison is clean. So this sync to head did not revert the Keras changes; it only rebased the same final file content onto current main so the PR is mergeable again.

[staticpayload]

Hi! With syncing to head you removed all changes that were previously added. Please sync to head carefully without reverting changes.

I messed up. It’s fixed now.

[RamSaw]

no commits pushed

[staticpayload]

To make this concrete: this was a force-push, not an additive commit on top of the old branch.

GitHub recorded the PR head moving from:

• c74f7956ed1a2307117ccabc535845297e096b1c

to:

• 944a6144bbf29fcfd394eb945dd35367f2396260

at 2026-04-01T13:35:58Z.

So the branch did move; it just moved via a history rewrite onto the latest main, which is why there is now a single replacement commit instead of an extra commit appended on top.

[ryan112358]

batch_selection.py is one example file where the changes at head are being reversed in this PR. There are more examples elsewhere, so please go back through and check everything carefully