flax

01854da1 - [No functional change] Delete VJPHiPrimitive and replace with HiPrim. This is a simple replace.

[No functional change] Delete VJPHiPrimitive and replace with HiPrim. This is a simple replace. If you are broken by this change, please rename `VJPHiPrimitive` to `HiPrim`. That's it! That's the change to unbreak. PiperOrigin-RevId: 968637756