Should we use a factored second moment estimator.

    Based on the shape of the variable.
    If we factor the accumulator, then this function returns a list of two
    mtf.Dimensions to reduce over.  We always pick the two largest dimensions.
    If there are not two dimensions of size >= min_dim_size_to_factor, then we
    do not factor.

    Args:
      shape: a Shape
    Returns:
      either a list of 2 Dimensions or None

def _factored_dims(self, shape):
    """Should we use a factored second moment estimator.

    Based on the shape of the variable.
    If we factor the accumulator, then this function returns a list of two
    mtf.Dimensions to reduce over.  We always pick the two largest dimensions.
    If there are not two dimensions of size >= min_dim_size_to_factor, then we
    do not factor.

    Args:
      shape: a Shape
    Returns:
      either a list of 2 Dimensions or None
    """
    if not self._factored or shape.ndims < 2:
      return None
    sorted_dims = sorted(shape.dims, key=lambda d: -d.size)
    if sorted_dims[1].size < self._min_dim_size_to_factor:
      return None
    return sorted_dims[:2]