Publishing with `py.typed`

[tmke8]

It seems most public-facing functions have type annotations in optax, so would the project be interested in publishing the package with a py.typed file for marking the package as having type annotations?

I can send a PR if there is interest.

[rdyro]

I think that's a great idea!

I need to catch up on how to verify this works correctly, do you have an idea on testing this?

[tmke8]

mypy and pyright on strict mode will complain if the py.typed file is not present, not sure about pytype. Perhaps one could also build the wheel, then uncompress it and check whether py.typed is there.

I looked at how haiku dealt with the file being added; it was in this PR: google-deepmind/dm-haiku#103 but it's not apparent if any testing was performed there.

[Shubhagata3499]

Hi! I’m interested in contributing and would like to work on this issue.

Here is how I plan to complete the task:

1. Add a py.typed file inside the main Optax package directory to mark the package as typed.

2. Update the package setup.cfg / MANIFEST.in (if required) to ensure py.typed is included when building wheels.

3. Build the wheel locally and verify that the py.typed file is included in the distribution.

4. Run local type check tests using tools like mypy or pyright in strict mode to confirm they detect the typing information.

5. Follow Optax's testing guidelines (e.g., running bash test.sh) to ensure no existing tests break.

Open a PR with a clear explanation of the changes, verification steps, and how it aligns with PEP 561 recommendations.