Load safetensors checkpoints per the target sharding in v1

Replace the transient-array multi-host load path in `SafetensorsLayout`
with a sharding-driven one: each process resolves which shards of each
tensor its own devices need, maps those index domains to byte ranges in
the file, reads exactly those ranges, and assembles the array with
`jax.make_array_from_single_device_arrays`. This drops the per-tensor
`jnp.sum` collective and the per-tensor `jax.jit` recompilation, and
collapses the single- and multi-process paths into one.

`index_domain_to_byte_runs` maps a shard's index domain to contiguous
byte runs -- one run for a leading-dimension shard, strided runs for an
inner-dimension shard -- and `coalesce` merges nearby runs into fewer
reads. Reads are issued as explicit ranged reads so a process transfers
only the bytes it needs.

The file index is a `_FileEntry` NamedTuple (`path`, `info`,
`data_start`) so call sites are self-describing; the no-`abstract_state`
branch lives in a module-level `_load_whole_tensors` helper so the
class's `_load` is purely the sharded path; and `_read_shard_bytes`
extracts requested runs from coalesced reads with a single forward sweep
through the merged blocks.

Expand the tests: byte-run math edge cases (non-zero offsets,
partial-both-dims, 3-D, non-4-byte itemsize), coalescing cases (overlap,
unsorted, exact-gap), `_read_shard_bytes` coalesced extraction, and
layout-level edge cases (shape mismatch, dtype cast, nested abstract
state, scalar, empty abstract state, duplicate tensor across files,
dtype sweep).

Author: mridul-sahu