Reduce Compilation Times for Numerical Hilbert Transform

[JoshPilipovsky]

Hello - I am working on a way to propagate distributions through neural networks (NNs) through using Characteristic functions (CFs). Long story short, in every layer of the NN, the hilbert transform needs to be applied to this function, which uses sinc and sine functions and sums them all up to compute the numerical hilbert transform. See code below as an example of a propagation step:

from jax import grad, jit, vmap
import numpy as np
import time
import matplotlib.pyplot as plt

def initializeGaussian(gridpoint, mu, sigma):
    CF_eval = jnp.exp(1j * mu * gridpoint - 0.5 * gridpoint ** 2 * sigma ** 2)
    return CF_eval

def CF_init(grid, mu, sigma, L, n):

    CF = jnp.zeros(shape=(L,n), dtype = np.complex_)
    for i in range(0, n):
        mui, sigmai = mu[i], sigma[i]
        CFi = vfunc_IG_jit(grid, mui, sigmai).block_until_ready()
        CF = CF.at[:, i].set(CFi)
    return CF

def hilbertTransform(f, grid, x, hilb_grid, h):
    sum = 0
    for m in hilb_grid:
        eval_pt = (x - m * h) / h
        sum += jnp.interp(m * h, grid, f) * jnp.sinc(eval_pt / 2) * jnp.sin(eval_pt / 2)
    return sum

def maxLayer(grid, CF, gridpoint, hilb_grid, h):
    CFz_eval = 0.5 * (1 + jnp.interp(gridpoint, grid, CF)) + 0.5 * 1j * (hilbertTransform(CF, grid, gridpoint, hilb_grid, h) - hilbertTransform(CF, grid, 0, hilb_grid, h))
    return CFz_eval

def CF_propagate(grid, CFx, L, n, hilb_grid):

    CFz = jnp.zeros(shape=(L,n), dtype = np.complex_)
    for i in range(0, n):
        CFxi = CFx.at[:, i].get()
        CFzi = vfunc_ML_jit(grid, CFxi, grid, hilb_grid, h).block_until_ready()
        CFz = CFz.at[:, i].set(CFzi)

    return CFz

# # # vectorize function maps # # #
vfunc_IG = vmap(initializeGaussian, in_axes = (0, None, None), out_axes = 0)
vfunc_ML = vmap(maxLayer, in_axes = (None, None, 0, None, None), out_axes = 0)

# # # jit it up # # #
vfunc_IG_jit = jit(vfunc_IG, static_argnums=(1,2))
vfunc_ML_jit = jit(vfunc_ML, static_argnums=(4))

# # # BEGINNING OF SCRIPT # # #

# define cutoff and resolution
d, L = 10, 100

# define HT resolution
h, M = 0.1, 50

# create grid along each axis
grid = jnp.linspace(-d, d, L)

# create grid for HT
hilb_grid = jnp.linspace(-M, M, 2 * M + 1)

# define initial Gaussian
n = 6
mu = np.random.uniform(-2, 2, n)
sigma = np.random.uniform(0.5, 4, n)

# warm-up
CFx = CF_init(grid, mu, sigma, L, n)
CFz = CF_propagate(grid, CFx, L, n, hilb_grid)

My issue is that this hilbert transform step takes a very long time to compile, especially for high-resolution grids (high values of M and L). I am trying to use vmap and jit to efficiently compile and run these propagations, but am not sure if I am doing this the correct way as I am very new to JAX. Essentially, I have an input gridpoint into the function maxLayer and that is being parallelized for all grid points. However, in the function hilbertTransform, I feel like the sum over all points in the hilb_grid is causing some slowdowns in the compilation. For large values of M, I sometimes get the following warning: E external/org_tensorflow/tensorflow/compiler/xla/service/slow_operation_alarm.cc:55]
Is there any way to parallelize that as well using vmap (kind of like a nested vmap)? Any thoughts would be appreciated!

[jakevdp]

Essentially any time you implement an algorithm with an explicit loop, you should expect compilation and execution times to be slower than necessary.

I think the main bottleneck here is your use of a loop within hilbertTransform. Instead, you should use broadcasting:

def hilbertTransform(f, grid, x, hilb_grid, h):
    eval_pt = (x - hilb_grid * h) / h
    return jnp.sum(jnp.interp(hilb_grid * h, grid, f) * jnp.sinc(eval_pt / 2) * jnp.sin(eval_pt / 2))

Also, you should delete the block_until_ready() calls in your functions, as they are preventing asynchronous dispatch from speeding up your loops.

I think your other loops could probably also be replaced with broadcasting operations as well, but since n is much smaller than m, the improvement will not be as dramatic.

[JoshPilipovsky]

I understand what you're saying and it makes sense. The broadcasting is a nice feature of python I didn't know about.

However, when I implement this fix, I am now getting the error TypeError: sub got incompatible shapes for broadcasting: (100,), (101,), in the eval_pt = (x - hilb_grid * h) / h step because the variable x is a batch dimension so it has length L while the hilb_grid has length 2M + 1. Correct me if I'm wrong but for broadcasting, we need x to be length 1.

[JoshPilipovsky]

Apologies, it was an error on my part. This solution stands - thanks once again