flaxって、何？

@Vengineerの戯言 : Twitter
SystemVerilogの世界へようこそ、すべては、SystemC v0.9公開から始まった 

Google Research から flax なるライブラリが

github.com

Flax: A neural network library for JAX designed for flexibility

ということなので、Google JAX のためのライブラリなんですね。

JAX なら TPU も使えるようなので。。。。

MNISTの例題もあります。

github.com