flax

9522bc0b - add static_argnums to nn.checkpoint