[pmap] In-line definitions of `jax.device_put_sharded` and `jax.device_put_replicated`.

Resolve unsoundness caught by pytype --strict-none-binding.

[pmap] Remove `jax.config.pmap_shmap_merge`.

Remove video wrapper to avoid directly calling ffmpeg.

Remove video wrapper to avoid directly calling ffmpeg.

Remove video wrapper to avoid directly calling ffmpeg.

[pmap] Handle edge cases in get_from_first_device for jax_pmap_shmap_merge.

Replace unicode escaped characters in ipynb files

[pmap] Avoid degraded performance under the new `jax.pmap`.

[acme] Remove `jax_pmap_shmap_merge=False`.

[acme] Explicitly set `jax_pmap_shmap_merge=False`.

Add an optional customizable loss callable in the PPONetworks, whose output will be added to total loss.

Silence pytype errors related to improved `jax.nn` type annotations.

Actually include all WPO files this time

Add Wasserstein Policy Optimization (WPO, http://arxiv.org/pdf/2505.00663) to Acme