#!/usr/bin/env python3

import math
import torch


def scale_to_bounds(x, lower_bound, upper_bound):
    """
    Scale the input data so that it lies in between the lower and upper bounds.

    Args:
        :attr:`x` (Tensor `n` or `b x n`):
            the input
        :attr:`lower_bound` (float)
        :attr:`upper_bound` (float)

    Returns:
        :obj:`torch.Tensor`
    """
    # Scale features so they fit inside grid bounds
    min_val = x.min()
    max_val = x.max()
    diff = max_val - min_val
    x = (x - min_val) * (0.95 * (upper_bound - lower_bound) / diff) + 0.95 * lower_bound
    return x


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

    Args:
        :attr:`train_inputs` (Tensor `n` or `n x d` or `b x n x d`):
            training data
        :attr:`ratio` (float, optional):
            Ratio - number of grid points to the amount of data (default: 1.)
        :attr:`kronecker_structure` (bool, default=True):
            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)

    Returns:
        :obj:`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)


def create_data_from_grid(grid):
    grid_size = grid.size(-2)
    grid_dim = grid.size(-1)
    grid_data = torch.zeros(int(pow(grid_size, grid_dim)), grid_dim, device=grid.device)
    prev_points = None
    for i in range(grid_dim):
        for j in range(grid_size):
            grid_data[j * grid_size ** i : (j + 1) * grid_size ** i, i].fill_(grid[j, i])
            if prev_points is not None:
                grid_data[j * grid_size ** i : (j + 1) * grid_size ** i, :i].copy_(prev_points)
        prev_points = grid_data[: grid_size ** (i + 1), : (i + 1)]

    return grid_data

@@ -42,16 +42,40 @@
 t torch%0A
+from typing import List%0A
 %0A%0Adef sc
@@ -1607,17 +1607,16 @@
  return 
-(
 ratio * 
@@ -1627,17 +1627,16 @@
 data
-)
 %0A%0A%0Adef c
 reat
@@ -1635,413 +1635,760 @@
 ef c
-reate_data_from_grid(grid):%0A    grid_size = grid.size(-2)%0A    grid_dim = grid.size(-1)%0A    grid_data = torch.zeros(int(pow(grid_size, grid_dim)), grid_dim, device=grid.device
+onvert_legacy_grid(grid: torch.Tensor) -%3E List%5Btorch.Tensor%5D:%0A    return %5Bgrid%5B:, i%5D for i in range(grid.size(-1))%5D%0A%0A%0Adef create_data_from_grid(grid: List%5Btorch.Tensor%5D) -%3E torch.Tensor:%0A    %22%22%22%0A    Args:%0A        :attr:%60grid%60 (List%5BTensor%5D)%0A            Each Tensor is a 1D set of increments for the grid in that dimension%0A    Returns:%0A        %60grid_data%60 (Tensor
 )%0A    
+ 
-prev_points = None%0A    for i in range(grid_dim):%0A        for j in range(grid_size):%0A            grid_data%5Bj * grid_size ** i : (j + 1) * grid_size ** i, i%5D.fill_(grid%5Bj, i%5D)%0A            if prev_points is not None:%0A       
+       Returns the set of points on the grid going by column-major order%0A            (due to legacy reasons).%0A    %22%22%22%0A    if torch.is_tensor(grid):%0A        grid = convert_legacy_grid(grid)%0A    ndims = len(grid)%0A    assert all(axis.dim() == 1 for axis in grid)%0A    projections = torch.meshgrid(*grid)%0A    grid_tensor = torch.stack(projections, axis=-1)%0A    # Note that if we did%0A
     
+#
     
@@ -2401,100 +2401,321 @@
 data
-%5Bj *
+ =
  grid_
-size ** i : (j + 1) * grid_size ** i, :i%5D.copy_(prev_points)%0A        prev_points =
+tensor.reshape(-1, ndims)%0A    # instead, we would be iterating through the points of our grid from the%0A    # last data dimension to the first data dimension. However, due to legacy%0A    # reasons, we need to iterate from the first data dimension to the last data%0A    # dimension when creating grid_data%0A   
  gri
@@ -2724,44 +2724,96 @@
 data
-%5B:
+ =
  grid_
-size ** (i + 1), : (i +
+tensor.permute(*(reversed(range(ndims + 1)))).reshape(ndims, -1).transpose(0,
  1)
-%5D%0A
 %0A   

Use meshgrid for create_data_from_grid