flax

03ea1a74 - Merge pull request #1629 from jheek:lifted-jacobian-products