flax

d2b105f0 - Migrate all usages of HiPrimitive to VJPHiPrimitive and delete HiPrimitive