Code for the paper Learning to optimize multigrid PDE solvers, which can be found on https://arxiv.org/abs/1902.10248.
The notation of the grid points is as follows: in any grid-related tensor, the (0,0) cell corresponds to the leftmost bottommost grid cell. The (I,J) cell then corresponds to the grid cell located in the I'th column and J'th row of the grid.