flax

a1543a81 - Light update to Flax NNX Transforms doc