Cast embeddings to a jnp array before slicing.
Cast embeddings to a jnp array before slicing.

This prevents an error when called with plain np params and a tracer for the ids, eg during an hk.dynamic_unroll.

PiperOrigin-RevId: 363608716
Change-Id: Ieda3a858f9093058a13c0d5043511c419ac5ae72

deepmind/dm-haiku,deepmind/dm-haiku