def should_precondition_dims(self):
        """A vector containing indicator indicating if the dim is preconditioned."""
        split_sizes = self._partitioner.split_sizes()
        rank = len(split_sizes)
        if self._preconditioner_type == PreconditionerType.ALL or rank <= 1:
            return [True] * rank
        else:
            return [True] * (rank - 1) + [False]
A vector containing indicator indicating if the dim is preconditioned.
tools/train/scalable_shampoo/distributed_shampoo.py