RET_CHECK failure: hlo->has_sharding() Side-effect HLO must have sharding

[sudhakarsingh27]

I'm facing following error while using pjit operator with multiple GPUs:

RuntimeError: Internal: RET_CHECK failure (external/org_tensorflow/tensorflow/compiler/xla/service/spmd/spmd_partitioner.cc:3779) hlo->has_sharding() Side-effect HLO must have sharding: %rng-get-and-update-state.2 = u64[2]{0} rng-get-and-update-state(), delta=1

Can someone help me understand what this is all about and how I can get past this?
Thanks.

Attached notebook (zipped) can be used to reproduce the error.
jax_pjit_error_repro.zip