Merge pull request google-deepmind#239 from michalsustr:master

PiperOrigin-RevId: 315790097
Change-Id: I9bf375a5e04dc0713be8f5c806c17e4349ecc808