flax

revert is_leaf logic in _check_carry_same_references

#4903

Merged