Add multi-device sharding for NSS and SwiG

[yallup]

Summary

Adds opt-in, single-host multi-device sharding to both BlackJAX NSS and the cache-aware Nested SwiG sampler introduced by #128.

• adds n_devices to BlackJAXNSSConfig and BlackJAXSwiGConfig, defaulting to the existing unsharded path
• shards live particles and replacement chains over a named JAX mesh while replicating the evidence integrator and adaptive covariance state
• uses one global likelihood all-gather for deletion and one full live-state all-gather for survivor starts per replacement
• keeps SwiG waveform caches ephemeral and device-local
• preserves the complete slice expansion, shrinkage, and acceptance diagnostics
• stores host-backed checkpoints and reshards state when resuming
• validates that n_live and n_delete are divisible by n_devices

This is a stacked PR targeting codex/nested-swig-cache. It should be retargeted to jim-dev after #128 merges.

Validation

• 221 passed, 3 skipped across tests/unit/samplers and the transient-likelihood unit suite
• 4 passed with XLA_FLAGS=--xla_force_host_platform_device_count=4, covering:
  • sharded NSS execution and diagnostics
  • sharded cache-aware SwiG likelihood consistency
  • host-backed checkpointing and resume
  • unavailable-device validation
• Ruff lint and formatting checks pass

Real multi-GPU performance and scaling still need validation on the target GPU node.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

🗂️ Base branches to auto review (1)

• .*-dev

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 5d448d7c-98dc-4155-b2fb-07477b9f4156

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch codex/nss-swig-sharding

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}