Add missing typing.Optional type annotations to function parameters.
Add missing typing.Optional type annotations to function parameters.

PiperOrigin-RevId: 376256736
Change-Id: I53e519a6651d1b47573cfd8efdddba9ee4e98ec6

deepmind/jraph