Skip to content

Navigation Menu

Appearance settings

Explore
By company size
By use case
By industry
View all solutions
Topics
- AI
- DevOps
- Security
- Software Development
- View all
Explore
- GitHub Sponsors
  Fund open source developers
- The ReadME Project
  GitHub community articles
Repositories
- Enterprise platform
  AI-powered developer platform
Available add-ons
Pricing

Search code, repositories, users, issues, pull requests...

Search

Clear

Search syntax tips

Provide feedback

We read every piece of feedback, and take your input very seriously.

Include my email address so I can be contacted

Saved searches

Use saved searches to filter your results more quickly

Name

Query

To see all available qualifiers, see our documentation.

Appearance settings

You signed in with another tab or window. Reload to refresh your session. You signed out in another tab or window. Reload to refresh your session. You switched accounts on another tab or window. Reload to refresh your session.

Dismiss alert

apple / axlearn Public

Notifications You must be signed in to change notification settings
Fork 369
Star 2.2k

Code
Issues 25
Pull requests 40
Actions
Security
Insights

Additional navigation options

Code
Issues
Pull requests
Actions
Security
Insights

Updated jax.tree_map to jax.tree_util.tree_map in update_transformation.py#1207

Open

m-ny wants to merge 4 commits intoapple:mainapple/axlearn:mainfrom

m-ny:mainm-ny/axlearn:mainCopy head branch name to clipboard

Conversation Commits4 (4)Checks Files changed

Open

Updated jax.tree_map to jax.tree_util.tree_map in update_transformation.py#1207
m-ny wants to merge 4 commits intoapple:mainfrom
m-ny:main

Commits

Commits on May 22, 2025

Update update_transformation.py
m-ny
authored
In base_layer_test.py: Updated import from jax.core to jax.extend.core Updated import from jax.interpreters.ad to jax.extend.ad Updated the function call from jax.interpreters.ad.deflinear to jax.e…
m-ny
authored
Updated import from jax.interpreters.pxla to jax.extend.pxla
m-ny
authored

Commits on May 26, 2025

Merge branch 'apple:main' into main
m-ny
authored

Footer

© 2025 GitHub, Inc.

Footer navigation

Terms
Privacy
Security
Status
Docs
Contact

You can’t perform that action at this time.