Confusion regarding traced objects and if-statements

[JoshPilipovsky]

Hi all - I am trying to optimize some of my code using vmap and jit but am running into some errors when the compiler looks at my jnp.where statement. Please see below.

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

# # # # # # # # # # # # # # # # # # # # # 

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

def intermediateLayer(grid, CFx, gridpoint, w, b):
    CFy_eval = jnp.exp(1j * b * gridpoint) * jnp.interp(gridpoint, grid, CFx)
    return CFy_eval

def computeHilbertTransform(f, grid, x, h, M):
    sum = 0
    hilb_grid = jnp.linspace(-M, M, 2 * M + 1)
    for m in hilb_grid:
        numerator = 1 - jnp.cos(jnp.pi * (x - m * h) / h)
        denominator = jnp.pi * (x - m * h) / h
        sum += jnp.interp(m * h, grid, f) * numerator / denominator
    return sum

def hilbertTransform(f, grid, x, h, M):
    hilb_grid_scaled = jnp.linspace(-M * h, M * h, 2 * M + 1)
    return jnp.where(x in hilb_grid_scaled, 0, computeHilbertTransform(f, grid, x, h, M))

def maxLayer(grid, CFy, gridpoint):
    h, M = 0.1, 100
    CFz_eval = 0.5 * (1 + jnp.interp(gridpoint, grid, CFy)) + 0.5 * (hilbertTransform(CFy, grid, gridpoint, h, M) - hilbertTransform(CFy, grid, 0, h, M))
    return CFz_eval

# # # # # # # # # # # # # # # # # # # # # 

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

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

# define initial Gaussian
mu, sigma = 0, 1

# define weight and bias
w, b = 0.5, 1

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

# # # jit it up # # #
vfunc_IG_jit = jit(vfunc_IG)
vfunc_IL_jit = jit(vfunc_IL)
vfunc_ML_jit = jit(vfunc_ML)

# warm-up
CFx = vfunc_IG_jit(grid, grid, mu, sigma).block_until_ready()
CFy = vfunc_IL_jit(grid, CFx, grid, w, b).block_until_ready()
CFz = vfunc_ML_jit(grid, CFy, grid).block_until_ready()

tic = time.perf_counter()
CFx = vfunc_IG_jit(grid, grid, mu, sigma).block_until_ready()
CFy = vfunc_IL_jit(grid, CFx, grid, w, b).block_until_ready()
CFz = vfunc_ML_jit(grid, CFy, grid).block_until_ready()
toc = time.perf_counter()
print(toc - tic)

# DEBUGGING - plot CFs
plt.plot(grid, CFx)
plt.plot(grid, CFy)
plt.plot(grid, CFz)
plt.show()

It is erroring on the line CFz = vfunc_ML_jit(grid, CFy, grid).block_until_ready() due to the following reason:

Traceback (most recent call last):
  File "/home/lqg/pythonUtils/main_vmap.py", line 65, in <module>
    CFz = vfunc_ML_jit(grid, CFy, grid).block_until_ready()
  File "/home/lqg/pythonUtils/main_vmap.py", line 35, in maxLayer
    CFz_eval = 0.5 * (1 + jnp.interp(gridpoint, grid, CFy)) + 0.5 * (hilbertTransform(CFy, grid, gridpoint, h, M) - hilbertTransform(CFy, grid, 0, h, M))
  File "/home/lqg/pythonUtils/main_vmap.py", line 31, in hilbertTransform
    return jnp.where(x in hilb_grid_scaled, 0, computeHilbertTransform(f, grid, x, h, M))
jax._src.errors.ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(bool[])>with<BatchTrace(level=1/1)> with
  val = Traced<ShapedArray(bool[100])>with<DynamicJaxprTrace(level=0/1)>
  batch_dim = 0
The problem arose with the `bool` function. 

See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.ConcretizationTypeError

If you look at the maxLayer function, then you can see that I used jnp.where to replace the if statement and that should fix the traced array error, but for some reason it doesn't. Any help would be appreciated!

[jakevdp]

Python's in operator cannot be overloaded, and so expressions like x in hilb_grid_scaled are not compatible with traced values. You could use something like jnp.any(x == hilb_grid_scaled) instead.

[JoshPilipovsky]

thanks for the prompt response

[jakevdp]

Actually, I just realized that my first statement is incorrect: in can be overloaded via the __contains__ method, so it would be possible to modify JAX such that your first attempt here would work. I'm going to think on that a bit...