flax

f80342a3 - Apply suggestions from code review

Apply suggestions from code review Co-authored-by: Marc van Zee <marcvanzee@gmail.com> Co-authored-by: Brennan Saeta <brennan.saeta@gmail.com>