Upstream patches needed: 1. https://github.com/FluxML/OneHotArrays.jl/pull/59 2. https://github.com/EnzymeAD/Reactant.jl/pull/1278