Is there a way to specify the pspec of a custom PyTree in shard_map?

[HeavyCrab]

I want to shmap a function with a custom PyTree arg (flax.struct.dataclass in fact). Therefore I need to specify the pspec of the arg. However the different fields of the arg have different pspecs. Is there a way to handle this?

[jakevdp]

You can pass a matching pytree with PartitionSpec instances as the leaves.