Add split_rng=False (current default) to sharded_map.
Add split_rng=False (current default) to sharded_map.

Haiku plans to make split_rng a required argument to hk.vmap in an upcoming
release. This change updates AlphaFold to preserve the current behaviour. We
also handle the case where users are using a release of Haiku without the
split_rng option, for these users split_rng=False is implied.

PiperOrigin-RevId: 428429472
Change-Id: I3292396eb330ffabf2f36f5364aa0ae1bc74cf71

deepmind/alphafold,deepmind/alphafold