Use Device hash and equality instead of using a (class, id) pair.
Use Device hash and equality instead of using a (class, id) pair.

We couldn't figure out why we did it this way in the first place and all the tests we have pass.

google/jax,google/jax,tensorflow/probability,tensorflow/probability,google/jax,google/jax