[README] nn libraries fix

[keren9898]

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

[mattjj]

Adding Keras sounds good, but why remove Equinox? That's a fantastic library, powering things like Stanford's Levanter (blog post). Maybe we should revise the 'Google X' part though...

[keren9898]

There are lot of great non-google jax software out there (including Equinox), but I think we can: 1) Keep the same rule for pointing to google projects (flax, keras, objax, ...) + point to open source list like awesome jax neural network section. 2) Change the rule by promoting some non-google open source software like Levanter with Equinox, but we mention some criteria to be featured.

On a side note, should we mention that grok is built using jax?

[patrick-kidger]

I'd like to keep Equinox there if possible, please :)

[jakevdp]

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.

[keren9898]

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.

[patrick-kidger]

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.

[jakevdp]

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.

[patrick-kidger]

Thanks Jake! Not disappointed with either you or Matt to be clear. :)

[keren9898]

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 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.

we need some criterion for why we're not including all the other great libraries mentioned in Awesome JAX and elsewhere.

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.

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.

We agree here. Providing a path for haiku users in the readme would have been prioritizing the jax community first.

[jakevdp]

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?

[patrick-kidger]

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.

[jakevdp]

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.

[patrick-kidger]

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 torch.nn, and the JAX version of torch.optim -- there's a reason PyTorch ships with those by default! So Equinox, Flax, Keras, Optax. The objective criteria is "is it one of those four". (A rule which has the added benefit of hopefully putting this long-standing topic to rest. :) ) We're no longer at the stage of the product lifecycle where that list is going to change, so thresholds like GitHub stars aren't going to be meaningful. (But we could still find one if that felt preferable to you.)

If the above seems reasonable then I would be happy to offer a pull request on this.

[jakevdp]

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.

[patrick-kidger]

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 torch.{nn,optim} equivalent) and (b) having a numerical threshold for determining inclusion.

[jakevdp]

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?

I suspect we have different approachs to curating such lists

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)

[patrick-kidger]

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"?

[jakevdp]

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?

[patrick-kidger]

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 dev branch for 2 months between releases!

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?

[jakevdp]

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!

[keren9898]

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.

Keras is great for new users.

As a practical level that is the JAX version of torch.nn, and the JAX version of torch.optim

Keras has these by default.

"any neural network or first-order optimisation library with at least 1k GitHub stars".

Why 1k ? Keras/Flax devs could come in and say make it 5k.

The objective criteria is "is it one of those four". (A rule which has the added benefit of hopefully putting this long-standing topic to rest. :) )

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?

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?

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.

[jakevdp]

I want to summarize quickly so we don't get lost in the details. As I see it, we have four options here:

1. Avoid listing/recommending downstream projects
2. List a small set of downstream projects, based on some non-objective criterion such as "what Jake likes"
3. List a small set of downstream projects, based on some yet-undetermined set of objective criteria
4. List all downstream projects that any maintainer requests to be listed

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).

[mattjj]

(2) SGTM, for some values of "Jake"

[gnecula]

@keren9898 I am curious, what is your interest here? Are you associated with any of these libraries?

[keren9898]

what is your interest here?

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.

Are you associated with any of these libraries?

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.

[jakevdp]

then Equinox support was dropped.

Can you elaborate? My impression is that Equinox is still very much an active project.

[lockwo]

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

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).

[lockwo]

• Avoid listing/recommending downstream projects
• List a small set of downstream projects, based on some non-objective criterion such as "what Jake likes"
• List a small set of downstream projects, based on some yet-undetermined set of objective criteria
• List all downstream projects that any maintainer requests to be listed

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).

[keren9898]

Can you elaborate? My impression is that Equinox is still very much an active project.

I meant no dedicated team support like keras or flax.

[keren9898]

P.S. penzai