Skip to content

fix: replace deprecated jax.tree_map with jax.tree.map#61

Merged
SimonDuToit merged 1 commit intoinstadeepai:mainfrom
shehper:fix/mypy-jax-tree-map-deprecation
Aug 4, 2025
Merged

fix: replace deprecated jax.tree_map with jax.tree.map#61
SimonDuToit merged 1 commit intoinstadeepai:mainfrom
shehper:fix/mypy-jax-tree-map-deprecation

Conversation

@shehper
Copy link
Contributor

@shehper shehper commented Jul 31, 2025

  • Fix mypy errors in mixer_test.py due to deprecated jax.tree_map usage
  • Replace 8 instances of jax.tree_map with jax.tree.map
  • Ensures compatibility with JAX >=0.4.25 and fixes development setup

This PR should resolve issue #60 .

The modified code passes mypy static type checks.

- Fix mypy errors in mixer_test.py due to deprecated jax.tree_map usage
- Replace 8 instances of jax.tree_map with jax.tree.map
- Ensures compatibility with JAX >=0.4.25 and fixes development setup

Fixes mypy static type checker errors that prevent pre-commit checks from passing.
@CLAassistant
Copy link

CLAassistant commented Jul 31, 2025

CLA assistant check
All committers have signed the CLA.

@SimonDuToit
Copy link
Collaborator

Hello, thanks for this fix!

Copy link
Collaborator

@SimonDuToit SimonDuToit left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good

@SimonDuToit SimonDuToit merged commit 20adc7a into instadeepai:main Aug 4, 2025
2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants