[pmap] Make clrs nets.py _get_first more robust under jax_pmap_shmap_merge

[copybara-service]

[pmap] Make clrs nets.py _get_first more robust under jax_pmap_shmap_merge

With jax_pmap_shmap_merge=True, the code in nets.py's NetChunked.call
was calling addressable_shards on inputs and hints without checking if they
are JAX arrays first. When numpy arrays are passed (e.g., during init), this
causes AttributeError.

This fix adds a robust _get_first helper that:

1. Checks if input is a JAX Array before accessing addressable_shards
2. For numpy arrays, uses simple x[0] indexing
3. For JAX arrays, handles 0-d arrays, SingleDeviceSharding, and properly
   extracts data from replicated shards