flax

97072989 - Merge pull request #5007 from pfackeldey:differentiable