Talk:JAX (software)

This article is rated C-class on Wikipedia's content assessment scale.
It is of interest to the following WikiProjects:

Software: Computing Low‑importance

	This article is within the scope of WikiProject Software, a collaborative effort to improve the coverage of software on Wikipedia. If you would like to participate, please visit the project page, where you can join the discussion and see a list of open tasks.SoftwareWikipedia:WikiProject SoftwareTemplate:WikiProject Softwaresoftware
Low	This article has been rated as Low-importance on the project's importance scale.
	This article is supported by WikiProject Computing.

Google Low‑importance

This article is within the scope of WikiProject Google, a collaborative effort to improve the coverage of Google and related topics on Wikipedia. If you would like to participate, please visit the project page, where you can join the discussion and see a list of open tasks.GoogleWikipedia:WikiProject GoogleTemplate:WikiProject GoogleGoogle

Low This article has been rated as Low-importance on the project's importance scale.

WikiProject Google To-do:

Here are some tasks awaiting attention:

Article requests : Articles for most of the other products listed here and here.
Assess : All articles in the Category:Unknown-importance Google articles and Category:Unassessed Google articles using the project's assessment scale
Expand : Google Mapathon, Google Talkback
Maintain : This WikiProject
Merge : Google Mobile Services into List of Google products
Stubs : Category:Stub-Class Google articles and Category:Google stubs
Update : List of features in Android and Gmail interface#Product integration. Update logos of Google Marketing Platform products
Other :
- Add more stuff to this to do list if you like! (click here...)
- create:
- Help the Google article for a good article status
- Improve the Outline of Google
- Get more members using :
{{subst:Wikipedia:WikiProject Google/Invite Members}}
- Infobox Images with transparent areas needing a different background color

Welcome

Hello, I created this page to have an accessible introduction to JAX, and give some MWEs of the 4 primary functions. I am also just learning JAX, so feel free to fix any mistakes. — Preceding unsigned comment added by Bquast (talk • contribs) 09:30, 21 June 2022 (UTC)[reply]

vmap example code

I found the vmap example code online (Apache 2.0 licenced). However, I am not quite sure how everything works. If anyone understands this example well and could add a line or two of explanation that would be great (or replace it altogether with a different simple example and an explanation of that example). Bquast (talk) 09:32, 21 June 2022 (UTC)[reply]

Google JAX

Pretty sure this is not the common name among practitioners, or developers. Never mind the disclaimer: "This is a research project, not an official Google product." Please consider moving this article elsewhere. 70.134.63.60 (talk) 07:23, 26 October 2022 (UTC)[reply]

What do you suggest we title the article? InfiniteNexus (talk) 03:36, 27 October 2022 (UTC)[reply]

@InfiniteNexus What about "JAX (Machine Learning framework)"? Otthorn (talk) 07:37, 31 January 2025 (UTC)[reply]

You can propose a move by following the instructions at WP:RM#CM. InfiniteNexus (talk) 07:45, 31 January 2025 (UTC)[reply]

Requested move 31 January 2025

The following is a closed discussion of a requested move. Please do not modify it. Subsequent comments should be made in a new section on the talk page. Editors desiring to contest the closing decision should consider a move review after discussing it on the closer's talk page. No further edits should be made to this discussion.

The result of the move request was: moved to JAX (software). Arbitrarily0 ^(talk) 14:02, 14 February 2025 (UTC)[reply]

Google JAX → JAX (Machine Learning framework) – If we look at all the papers on arXiv containing "JAX" in their name in the category CS (55 papers), only one use "Google JAX" in it's abstract and another one use "Google's JAX". All the 53 others will cite it has "JAX" like in "EagerPy is a Python framework that lets you write code that automatically works natively with PyTorch, TensorFlow, JAX, and NumPy." for example.

No mention of "Google JAX" in the Github README.

No mention of "Google JAX" on the front page of the documentation. Not even a mention of "Google". It is not even an official Google product: "This is a research project, not an official Google product" from the Github page.

No mention of "Google JAX" in the original paper. Otthorn (talk) 08:24, 31 January 2025 (UTC) — Relisting. ~/Bunnypranav:<ping> 13:07, 7 February 2025 (UTC)[reply]

Edit: Per the below discussion, I'd prefer moving to JAX (software) instead. Malerisch (talk) 12:13, 13 February 2025 (UTC)[reply]

Support a move to JAX (machine learning framework): the disambiguator should be lowercased since it isn't a proper noun. Malerisch (talk) 16:16, 12 February 2025 (UTC)[reply]
Jax isn't a full machine learning framework; all it does is autodiff. Maybe "JAX (Autodifferentiation Framework)"? Pvlkmrv (talk) 18:50, 12 February 2025 (UTC)[reply]
I actually hate the word "framework" for these things. Maybe "library"? Pvlkmrv (talk) 18:51, 12 February 2025 (UTC)[reply]
On both the Github and the frontpage they use the term library

"JAX is a Python library for accelerator-oriented array computation and program transformation, designed for high-performance numerical computing and large-scale machine learning."

So I would agree that library is a good choice here. Otthorn (talk) 18:53, 12 February 2025 (UTC)[reply]
Alternative disambiguators could be "(software)", as in Caffe (software), Theano (software), and pandas (software), or "(machine learning)", as in Torch (machine learning). "(library)" is too vague; some may think that it refers to a physical library. Malerisch (talk) 19:17, 12 February 2025 (UTC)[reply]

Is "JAX" an acronym? Is there a reason to use "JAX" instead of "Jax"? —⁠ ⁠BarrelProof (talk) 01:09, 14 February 2025 (UTC)[reply]
According to the original JAX paper [1] and this discussion [2], it was originally an acronym for "Just After eXecution" but is now just known as "JAX". Former acronyms/initialisms that became the official name, such as SAT, KFC, AARP, BP, etc., are still written in all-caps. Malerisch (talk) 01:28, 14 February 2025 (UTC)[reply]
Support Hplotter (talk) 10:09, 13 February 2025 (UTC)[reply]
Is there a way to change the proposed title? Or should we just wait for the discussion to be closed and then the person changing the page name will take the consensus for the final edit?

I think @Malerisch proposition makes the most sense. Otthorn (talk) 10:29, 13 February 2025 (UTC)[reply]
What is your preferred move target? I proposed a few options. I personally think JAX (software) may be the best choice: it's more WP:CONCISE than JAX (machine learning framework), and the "(machine learning)" disambiguator seems like it's used more often with machine learning concepts, like neural network (machine learning), attention (machine learning), or hyperparameter (machine learning), than with software libraries. Malerisch (talk) 11:34, 13 February 2025 (UTC)[reply]
With a quick intitle search I would prefer the Jax (software) option. "(machine learning)" does seem to refer to concept more than libraries, "(library)" does not seem relevant here.

For "(machine learning framework") I only found Flux (machine-learning framework), but the title is also over hyphenated in my opinion so it's probably a good example, and it's a single page. Otthorn (talk) 11:57, 13 February 2025 (UTC)[reply]
probably not* a good example Otthorn (talk) 11:58, 13 February 2025 (UTC)[reply]
Torch (machine learning) and Flux (machine-learning framework) need a more specific disambiguator than "(software)" since it's not WP:PRECISE enough (see Torch (disambiguation)#Computing and technology and Flux (disambiguation)#Computing); this isn't the case with JAX however. I think you can modify the proposed title, although your comment should be enough for the closer. Malerisch (talk) 12:13, 13 February 2025 (UTC)[reply]

The discussion above is closed. Please do not modify it. Subsequent comments should be made on the appropriate discussion page. No further edits should be made to this discussion.

Removal of examples?

I don't think code examples should be on wiki. It's not meant as a doc/tutorial for framework. Perhaps description of jax's features would be sufficient (like the tensorflow page). Features like just in time compilation (jit), support for SPMD, distributed arrays and automatic parallelization, etc. — Preceding unsigned comment added by Anothercat613 (talk • contribs) 12:02, 26 March 2025 (UTC)[reply]

I've started work on refactoring the page to remove technical details and replace them with plain English explanation . If there's no issue, I'll remove the code examples soon. Anothercat613 (talk) 08:53, 6 April 2025 (UTC)[reply]