def __init__(
        self,
        param,
        block_size,
        merge_small_dims_block_size,
        best_effort_shape_interpretation,
        preconditioner_type=PreconditionerType.ALL,
    ):
        """Initializes the preconditioner.

        Args:
          param: parameter to precondition.
          block_size: Block size used to split param.
          merge_small_dims_block_size: Block size for merging dims.
          best_effort_shape_interpretation: Whether to collapse/merge dims together.
          preconditioner_type: Type of preconditioner to use.
        """
        self._original_shape = param.shape
        self._transformed_shape = param.shape
        if best_effort_shape_interpretation:
            self._transformed_shape = merge_small_dims(
                self._original_shape, merge_small_dims_block_size
            )
        reshaped_param = jnp.reshape(param, self._transformed_shape)
        self._partitioner = BlockPartitioner(reshaped_param, block_size)
        self._preconditioner_type = preconditioner_type
Initializes the preconditioner.

        Args:
          param: parameter to precondition.
          block_size: Block size used to split param.
          merge_small_dims_block_size: Block size for merging dims.
          best_effort_shape_interpretation: Whether to collapse/merge dims together.
          preconditioner_type: Type of preconditioner to use.
        
tools/train/scalable_shampoo/distributed_shampoo.py