Support mixed SingleReplicaArrayRestoreArgs / ArrayRestoreArgs batches in SingleReplicaArrayHandler.deserialize

[anveshjayakrishna]

Summary

SingleReplicaArrayHandler.deserialize currently requires every arg in the batch to be SingleReplicaArrayRestoreArgs, rejecting anything else:

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 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 (each rank reads independently).
• reassemble results in the original arg order.

Note that SingleReplicaArrayRestoreArgs subclasses ArrayRestoreArgs, 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)?

[cpgaffney1]

This is simple enough, I think there is no problem with this.

The only other way to do it currently would be to provide a custom TypeHandlerRegistry, but that would be rather cumbersome compared to this proposal.
https://orbax.readthedocs.io/en/latest/api_reference/checkpoint.handlers.html#pytreecheckpointhandler