Framework agnostic layers

[cgarciae]

Recently Jraph came out with some functional framework agnostic layers and that got me thinking that maybe it would be nice if as a community we could have some functional framework-independent layers that all frameworks could build their framework specific layers on-top.

This could include:

• Parameter initializers
• Pure functional layers
• Regularizers
• Losses
• Metrics
• etc

Basically all the building blocks you need to create a DL framework, the code would be well tested if adopted, and frameworks could concentrate on providing better APIs instead.

Given the relation between the flax and jax team I am posting this here but I don't mean to imply that this has to live in the flax codebase, it could be its separate repo. In Elegy we have a bunch of functional versions of losses and metrics ported from Keras we could contribute.

[jheek]

Too a large extend I think the JAX ecosystem is already moving in this direction. For example activation functions and initializers used to be part of every NN library but now they are part of jax.nn. Another great example is CLU which helps with metric reporting and other common training loop utilities.

In general I think that because JAX is functional we can make libraries integrate well by making sure their APIs are functional too. For example Flax has a pretty basic Optimizer module but you can also use the optax interchangeably.

The flipside is that we want to avoid a situation where you need to have a dozen dependencies before you can build anything. So we sometimes end up with a compromise: A loosely coupled module that works beyond Flax modules but is still included in the same package to avoid fragmentation.

As the JAX ecosystem matures I think more single-purpose libraries will mature and we will end up deprecating more in-house sub modules.