Support mixed SingleReplicaArrayRestoreArgs / ArrayRestoreArgs batches in SingleReplicaArrayHandler.deserialize

## Summary

[`SingleReplicaArrayHandler.deserialize`](https://github.com/google/orbax/blob/main/checkpoint/orbax/checkpoint/_src/serialization/jax_array_handlers.py#L1901) currently requires *every* arg in the batch to be `SingleReplicaArrayRestoreArgs`, rejecting anything else:

```python
for arg in args:
if not isinstance(arg, SingleReplicaArrayRestoreArgs):
raise ValueError(
'Must provide `SingleReplicaArrayRestoreArgs`, but got'
f' {type(arg)}.'
)
```

This makes single-replica restore all-or-nothing. Once the handler is registered for `jax.Array`, every array in the pytree must go through the single-replica broadcast path.

## Motivation

Some arrays cannot be correctly broadcast from a single replica — most notably arrays **partitioned along the replica axis**, which hold distinct data per replica. Broadcasting one replica's shard to all others would corrupt them; they need a plain [`ArrayRestoreArgs`](https://github.com/google/orbax/blob/main/checkpoint/orbax/checkpoint/_src/serialization/jax_array_restore_args.py) so each rank reads its own shard.

Today there is no way to opt individual arrays out of the broadcast path within a single restore, so a checkpoint that mixes broadcast-eligible and replica-partitioned arrays cannot use single-replica restore at all.

## Proposal

Split the batch by arg type instead of rejecting:

- `SingleReplicaArrayRestoreArgs` → existing single-replica broadcast path.
- plain `ArrayRestoreArgs` → delegate to the parent [`ArrayHandler.deserialize`](https://github.com/google/orbax/blob/main/checkpoint/orbax/checkpoint/_src/serialization/jax_array_handlers.py) (each rank reads independently).
- reassemble results in the original arg order.

Note that [`SingleReplicaArrayRestoreArgs` subclasses `ArrayRestoreArgs`](https://github.com/google/orbax/blob/main/checkpoint/orbax/checkpoint/_src/serialization/jax_array_restore_args.py), so the split must test `SingleReplicaArrayRestoreArgs` **first**.

## Backwards compatibility

When all args are `SingleReplicaArrayRestoreArgs` (today's usage), the standard-arg group is empty, the parent handler is never invoked, and the original broadcast path runs unchanged.

## Correctness note

The broadcast path is a collective, so all ranks must agree on which arrays are single-replica vs. standard — callers must derive args from the same sharding tree on every rank.

## Question for maintainers

Would you accept extending `deserialize` to handle mixed batches this way, or was the strict all-`SingleReplicaArrayRestoreArgs` requirement intentional (e.g. a dispatcher/broadcast constraint I'm missing)?

Start at SingleReplicaArrayHandler.deserialize in orbax/checkpoint/_src/serialization/jax_array_handlers.py and read the related ArrayHandler.deserialize implementation. Then inspect SingleReplicaArrayRestoreArgs and ArrayRestoreArgs in jax_array_restore_args.py, including their inheritance. Done means mixed batches use the appropriate path, results retain original order, and all-single-replica batches keep the existing behavior.

Written by the indexing model from the issue text.