New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[README] nn libraries fix #20745
Comments
There are lot of great non-google jax software out there (including Equinox), but I think we can:
On a side note, should we mention that grok is built using jax? |
I'd like to keep Equinox there if possible, please :) |
In the past we have tried to come up with objective criteria for which libraries to include and not include in sections like this, and have not been able to come up with any. I believe the current list was created with the loose criteria "projects managed by google", in which case the current PR is the correct change (with apologies to Equinox, which unless I'm mistaken is no longer actively developed by googlers). That said, dropping Equinox doesn't sound right to me. But if we keep Equinox, we need some criterion for why we're not including all the other great libraries mentioned in Awesome JAX and elsewhere. Absent some good criteria to choose what we include, the easiest fix may be to drop this section altogether. These sorts of complications are why we've generally avoided adding lists of downstream projects to the JAX docs. |
For context, I believe the criteria for 1) is the same criteria used by @patrick-kidger in his PR to Remove Haiku and add Equinox in the readme. I think that PR should have clarified a path for haiku users (i.e. move to flax as recommended by Haiku) instead of removing the Haiku link completely. I believe that the Flax and Haiku teams have guides on how to transition. |
I'm afraid not. Haiku was removed as the library is deprecated for greenfield projects. I believe the bar has never been "Alphabet-affiliated", and indeed I recall from internal discussions with the JAX team that they have worked hard for the project not to feel like it's Google-only. This part of the JAX README serves as an entry point for new users, and indeed Equinox has seen a significant uptick in its rate of new users since it was added here. I think this is evidence that this section of the README is exceptionally useful for the community. I'm a little disappointed that I seem to need to defend the inclusion of my library -- which seems to have been very useful to many -- so strongly. |
I’m sorry you’re feeling disappointed Patrick, but I hope you can take solace in the fact that the two JAX maintainers who have weighed in in this thread have also advocated for keeping Equinox on the list. |
Thanks Jake! Not disappointed with either you or Matt to be clear. :) |
I am sorry you feel this. You have done undeniably good work. Putting the jax community first, my concern is with the continuation of Equinox given that its maintained only by you and not by google anymore.
Giving equinox an edge, will also give other patrick's work an edge over other open-source work, and will likely inhibit new ideas in the jax community and/or create a sense of unfairness. One example is lineax, PR vs cola among others.
We agree here. Providing a path for haiku users in the readme would have been prioritizing the jax community first. |
I propose that absent a good objective means to choose which libraries to include and exclude here, we remove the section entirely. That will be fair to every downstream package. What do you think? |
I think this would be a worst possible outcome. As above, I think this provides an exceptionally valuable onboarding resource for new users. I would like to argue against this in the strongest possible terms. |
OK, then do you have a proposal for an objective criterion we can use to choose which libraries to include here, and avoid accusations of being unfair to the projects we don't include? Absent that, an alternative approach is a free-for-all, where we'd include any library requested to be listed here, to avoid disappointing the maintainer. I think that's probably a worse outcome than removing this section entirely. |
Sure. So I think this section is basically "what, besides JAX, does a new user need in order to get started?" The audience are specifically new users. As a practical level that is the JAX version of If the above seems reasonable then I would be happy to offer a pull request on this. |
I think "your library is not one of the blessed four" is not the kind of objective rule that will successfully avoid potential accusations of unfairness... For example, I suspect that if I declared that Flax, Keras, and Optax were the three libraries that cover relevant functionality, you may be disappointed when I tell you that the reason we didn't include Equinox was because it did not meet the objective criterion of being on that list. |
I suspect we have different approachs to curating such lists :) When I curate lists of this sort (here's a non-JAX example), the objective criteria has been things like "have I enjoyed it". Not all such criteria need be numerical! But anyway, if you prefer something numerical, then as per the second half of my message -- "any neural network or first-order optimisation library with at least 1k GitHub stars". This meets the requirements of being (a) of immediate interest to new users (i.e. a |
I did consider using github stars as a criterion, but there are many JAX-related libraries and repositories out there that meet the criterion, so it's not that simple. For example, I don't think we'd want to list Haiku or Trax, despite them easily meeting your criteria. As for "have I enjoyed it", I can say I've not used Equinox or Keras much compared to Flax and Haiku, so by that rubric I'd personally have to lean toward leaving both of them out. Does that sound reasonable to you?
There's a difference between a list maintained by an individual, and a list maintained by a community. If it were just me, I'd have already made a decision (and possibly one you would object to) |
I considered adding "is the library maintained for greenfield projects" to the criterion as well, but I took that bit as a given! (AFAIK this excludes Trax.) Indeed I'm not suggesting that on this occasion we make this decision based on what you, personally, have happened to experience. My example is meant to highlight that non-numerical criteria may still be both useful and "objective". So: "any neural network or first-order optimisation library, that is still maintained for greenfield projects, with at least 1k GitHub stars"? |
Yeah, that's getting there I think. What do you mean when you say "that is still maintained"? For example neural_tangents arguably meets all of your proposed criteria, but its most recent commit was 2 months ago. Does it fit in the list? If not, do we need to add it to the list if the maintainers update the README someday? |
I wouldn't measure it by 'most recent commit' -- for the reason you highlight; because we should leave space for a library to be considered mature; and because development may simply be occuring out-of-sight on a I think it's a measure of reliability: how engaged are the maintainers with the community; are serious bugs fixed in a timely fashion. How might you formalise that, do you think? |
I don't know – I've tried and failed to come up with such objective criteria in the past, which is why I'm asking questions about yours! |
Keras is great for new users.
Keras has these by default.
Why 1k ? Keras/Flax devs could come in and say make it 5k.
How would you feel if the jax team had an objective criteria "is it one of those four", but instead of Equinox it was Haiku?. Would you be encouraged to start your library?
Experienced jax users could point that Equinox stateful operation (e.g. for kv cache, batchnorm, ...) is based around a jax bug. This would break Equinox and downstream libraries (e.g. Levanter) the moment that bug is fixed. As a matter of fact, This happened before when equinox used to monkey patch parts of jax that got changed. This design would be considered as a serious bug, and would mark Equinox unsafe for new users. By this standard, we can exclude equinox. Overall, I do not think gate keeping is useful in any context here, this is why I suggested 1), that offers new jax users with a selection of libraries backed by dedicated google teams, and for more curious/advanced users with an awesome list for more libraries out there. We can of course remove that section entirely as @jakevdp suggested, and let each library promote itself as they wish, if we can not reach a solution here. |
I want to summarize quickly so we don't get lost in the details. As I see it, we have four options here:
Historically we've typically chosen (1), because anything else leads to threads like this one. This README section is an exception to that, which has mainly persisted due to flying a bit under the radar. Patrick has said above that (1) is the worst possible outcome. Patrick is advocating for (2), but I think that opens us up to future troubles (as evidence, I'd point to the strong reactions present in this thread any time one of the participants has gotten a whiff that their favorite library might be excluded). (3) would be ideal, aside from the small problem that such objective criteria are difficult to come up with. (4) in my mind is the worst possible outcome, because hitting beginners with an unbounded list of possibilities is arguably less useful than saying nothing. If (3) is not possible, I'd probably opt for either (1) (at the expense of JAX users, but letting me personally keep some friends) or (2) (to the benefit of JAX users, but causing me personally to lose some friends). |
(2) SGTM, for some values of "Jake" |
@keren9898 I am curious, what is your interest here? Are you associated with any of these libraries? |
I raised this issue due to my frustration with the stability of the JAX neural network libraries. Initially, I began using Haiku, but it was suddenly discontinued. I then switched to Equinox. but then Equinox support was dropped. Now, observing Flax, it seems likely to undergo changes again. This is why I suggested adding the more stable Keras API.
No. I am an early JAX adopter (and autograd/tf-keras user before that). Having said that, All libraries have their issues, but at least to be suggested for new or experienced users, we need a framework that has a stable API, free from bugs, and maintained by some team. I admire Haiku,Equinox,Flax, and Keras hard work. But we should always prioritize the jax user base. Lastly, as @patrick-kidger noted, being featured in the README - and else where in JAX repo- provides a significant visibility boost. However, I am concerned that offering this boost to a single developer without some criteria could be unfair to downstream projects and JAX users who anticipate some level of long-term support. I also think Patrick setting criteria that seem tailor-made to include his own work, while potentially gatekeeping other projects (including other Google projects), does not send the right message. |
Can you elaborate? My impression is that Equinox is still very much an active project. |
I am also interested in this, what jax bug does this rely on? My understanding was how equinox does stateful operations is not too different from other functional approach (e.g. the State monad in Haskell, https://en.wikibooks.org/wiki/Haskell/Understanding_monads/State). |
I agree (1) and (4) should not the be target. To me, (2) and (3) are not all that different. You/google/jax team could create some set of criteria (stars, last commit, activity, etc.) that are "objective" but still entirely arbitrary, these things are being pulled out of thin air and even if they refer to objective quantities, it doesn't seem effectively that different from the BDFL approach of (2). |
Description
Small update: Add Keras to the README neural network section, and remove equinox @jakevdp @fchollet @patrick-kidger
System info (python version, jaxlib version, accelerator, etc.)
NA
The text was updated successfully, but these errors were encountered: