flax

a4dade8c - Use ml_collections for hyperparameter handling