import math


def choose_grid_size(train_inputs, ratio=1.0, kronecker_structure=True):
    """
    Given some training inputs, determine a good grid size for KISS-GP.

    :param x: the input data
    :type x: torch.Tensor (... x n x d)
    :param ratio: Amount of grid points per data point (default: 1.)
    :type ratio: float, optional
    :param kronecker_structure: Whether or not the model will use Kronecker structure in the grid
        (set to True unless there is an additive or product decomposition in the prior)
    :type kronecker_structure: bool, optional
    :return: Grid size
    :rtype: int
    """
    # Scale features so they fit inside grid bounds
    num_data = train_inputs.numel() if train_inputs.dim() == 1 else train_inputs.size(-2)
    num_dim = 1 if train_inputs.dim() == 1 else train_inputs.size(-1)
    if kronecker_structure:
        return int(ratio * math.pow(num_data, 1.0 / num_dim))
    else:
        return ratio * num_data