Rename 'master' to 'main' in the JAX APIs.
Rename 'master' to 'main' in the JAX APIs.

PiperOrigin-RevId: 329167360
Change-Id: If611497a8bd2d3a5fbc08e591d84b291fc9667a1

deepmind/dm-haiku,deepmind/dm-haiku