def _remove_leading_sharding_annotation(pspec):
        """Mapping from N-d to (N-1)-d, used for quantization, factoring etc."""
        # None and PSpec(None) are valid PSpecs.
        if pspec and len(pspec) > 1:
            return pjit.PartitionSpec(*pspec[1:])
        else:
            return []
Mapping from N-d to (N-1)-d, used for quantization, factoring etc.
tools/train/scalable_shampoo/distributed_shampoo.py