Hello, I am trying to implement a “custom” convolution routine to perform spherical convolutions.

Basically, I store my spherical image as a 1D array following the healpix pixelisation scheme. I am trying to optimize an already existing routine and I would like to re-implement the convolution step to possibly make it more efficient.

A bit of context, you can skip to the last paragraph for the actual question.

For each element of the input array, I have a function which lets me compute the index of the neighboring pixels. The idea is to perform a convolution on the neighboring pixels of each element of the array. This idea works pretty well, but I fear that the implementation is not optimal.

At the moment, for each pixel I am gathering the neighboring pixels and parsing them in a new 1D array (using the `tf.gather`

) , which results to be a 1D array 9 times bigger on which we perform the convolutions (since each pixel has ~8 neighboring pixels). Once I have such an “ordered array” I can perform simple 1D convolutions with stride=9 to achieve a spherical convolution. For more information about this idea, please take a look at the original paper: [1902.04083] Convolutional Neural Networks on the HEALPix sphere: a pixel-based algorithm and its application to CMB data analysis

The actual question:

What I would like to do now is to tell tensorflow to perform a “custom” convolution, given the index of the elements that I want to convolve. Basically I would like to give tensorflow a series of indexes (one set of indexes for each element of the 1D array) and tell it to convolve them directly. This would involve accessing elements from the 1D array not contiguously, I think.

For example, for element 0 of the array, the convolution would involve the elements [ 0, 4, 11, 3, 2, 1, 6, 5, 13], for pixel 1 [1, 6, 5, 0, 3, 2, 8, 7, 16], for pixel 2 [ 2, 8, 7, 1, 0, 3, 10, 9, 19] and so on.

Is there a way to achieve it? Thank you very much for your help!