Small cleanup of DistributedLayout logic. No functional changes.
Small cleanup of DistributedLayout logic. No functional changes.

PiperOrigin-RevId: 439269183
Change-Id: I96e27c63c09f4d6d49ecad5e66e39c10a8e6e8a9

deepmind/acme,deepmind/acme