flax

6844a540 - Merge pull request #3741 from chiamp:einsum_nnx