flax

22468a1b - Allow DenyList to be compared.

Allow DenyList to be compared. This fixes e.g. the use of DenyList as a dict key in a nn.while_loop split_rngs arg, as JAX sorts dict keys in tree operations. PiperOrigin-RevId: 881082936