feat: add Type hints for examples/dp_logistic_regression.py and examples/jax_api_example.py

[Neerajpathak07]

Progresses #124

This PR adds a type hint for the examples file dp_logistic_regression.py.
Optimizing the function from :-

def logistic_loss(params, feature_matrix, labels):

to:-

def logistic_loss(
    params: Mapping[str, jax.Array],
    feature_matrix: jax.Array,
    labels: jax.Array,
) -> jax.Array:

similarly for other functions throughout the file.

[Neerajpathak07]

@ryan112358 the CI lint error as of now is failing for pyink I am awaiting on the pyink issues to be fixed for the test files which should be done when #90 merges in after which I'll pull it here.

[ryan112358]

Please rebase to head, once you pass the CI checks, we can move onto internal review and merge

[Neerajpathak07]

@ryan112358 Yeah I was waiting for the Pyink formatting for accounting.py to be added. I'll do it now.

[codecov-commenter]

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

Codecov Report

✅ All modified and coverable lines are covered by tests.
⚠️ Please upload report for BASE (main@b3b6937). Learn more about missing BASE report.
❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #130   +/-   ##
=======================================
  Coverage        ?   71.67%           
=======================================
  Files           ?       25           
  Lines           ?     2655           
  Branches        ?        0           
=======================================
  Hits            ?     1903           
  Misses          ?      752           
  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.

[Neerajpathak07]

@ryan112358 Have resolved the CI issues. Now it's only up to the copybara check which If I'm not wrong has to be done internally by you guys.

[ryan112358]

Hi Neeraj,

Our internal checks show an invalid import order (typing should be moved up to the standard library section). To fix, please run "isort" on your changed files. Bonus would be if you can also update .github/workflows/ci.yml to include an "isort --check" action while you are at it so that this can be caught earlier in the future

[Neerajpathak07]

@ryan112358 CI does fail at the top level for files included in tests/, examples/ and jax-privacy/. There 2 things we can do:-

1. change the CI to only test for examples (won't recommend as per CI rule present upstream)
2. fix all the upstream files by running isort (45 files to be changed in total)

LMK what would you opt for.

[ryan112358]

@ryan112358 CI does fail at the top level for files included in tests/, examples/ and jax-privacy/. There 2 things we can do:-

1. change the CI to only test for examples (won't recommend as per CI rule present upstream)
2. fix all the upstream files by running isort (45 files to be changed in total)

LMK what would you opt for.

Thanks for checking, I guess there is some difference in our internal config for isort, so changing the files will break things internally. For now, just revert the the change to ci.yml and manually make the requested change (moving the typing import to the top-level block of imports, in alphabetical order). We will try to figure out internally how we can update ci.yml to be consistent between internal and external import checks.

[Neerajpathak07]

@ryan112358 CI does fail at the top level for files included in tests/, examples/ and jax-privacy/. There 2 things we can do:-

1. change the CI to only test for examples (won't recommend as per CI rule present upstream)
2. fix all the upstream files by running isort (45 files to be changed in total)

LMK what would you opt for.

Thanks for checking, I guess there is some difference in our internal config for isort, so changing the files will break things internally. For now, just revert the the change to ci.yml and manually make the requested change (moving the typing import to the top-level block of imports, in alphabetical order). We will try to figure out internally how we can update ci.yml to be consistent between internal and external import checks.

works

[Neerajpathak07]

@ryan112358 wondering if the internal linter has something to say about the import order again. Given that this was fixed prev by isort itself.

[ryan112358]

@ryan112358 wondering if the internal linter has something to say about the import order again. Given that this was fixed prev by isort itself.

Still an issue. If you don't change the order or structure of existing imports, the internal check should pass. So please revert the recording to that. All you have to do is put the new imports in the correct locations. The generally structure is

1. Top level block: standard lib imports (like typing)
2. Middle block: standard third party imports (Jax, numpy, etc.)
3. Final block: internal imports within files from the jax-privacy library

[Neerajpathak07]

Before adding a change intending to order the imports this way:-

from typing import Any, Mapping, Tuple

from absl import app
import jax
import jax.numpy as jnp
import jax_privacy
from jax_privacy import batch_selection
from jax_privacy.experimental import execution_plan
import numpy as np
import optax

just adding the new line at the top of the old order!!

[Neerajpathak07]

@ryan112358 Thank you for your follow-ups!!
Just need a little heads-up on a few things- could there be a way possible for contributors to get an idea where the internal checks fail as when pushing a commit which runs clean locally and on the CI but fails internally it creates a little confusion as to what went wrong and increase review-rounds back and forth causing extra hassel!!!

Maybe there could be a better way to make these contributions seamless?

[Neerajpathak07]

NOTE:- Once this PR lands in it should fix the error caused for the CI unit tests for all the open PR's by contributors.
Ref:- https://github.com/google-deepmind/jax_privacy/actions/runs/22150985293/job/64041979147?pr=131