Update deps pinned to jax version 0.8.2

[dipannita08]

Description

Background: A previous PR upgrading google-cloud-mldiagnostics through seed-env resulted in dependency issues caused by (1) cloud-tpu-diagnostics version 1.1.14 being yanked but dependency resolver unable to detect and exclude this (b/477743392) and (2) TPU lock file unable to pick up version google-cloud-mldiagnostics 0.5.10. This PR also inadvertently removed jaraco-functools, causing some more CI failures and was rolled back.

This change is attempting do a full update of both TPU and GPU deps.

FIXES: b/481140947

Tests

• Run MaxText local
• CI Tests

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[bvandermoon]

It's surprising to me that the GPU tests are failing when the GPU generated requirements file hasn't changed. Do you have any thoughts on why that would be?

[dipannita08]

It's surprising to me that the GPU tests are failing when the GPU generated requirements file hasn't changed. Do you have any thoughts on why that would be?

Yeah, I'm not sure either. Ran out of bandwidth last week to dig deeper.

[kanglant]

Why are the jaxlib version and the jax version mismatched?

[dipannita08]

Good question - these are moved from seed-env generated files - I'm not sure why this is the case.

[kanglant]

Which jax seed commit are you using?