jax

812d2462 - don't require passing identity value. It isn't the initial value - identity is required for implementation correctness