vmap gives wrong values when applied to a piecewise linear function.

[idnm]

I found this strange behavior of vmap which looks like a bug to me. Here is the code defining a piecewise linear function.

import jax.numpy as jnp
from jax import vmap
import matplotlib.pyplot as plt

def line(x, x0, y0, x1, y1):
    return (y1-y0)/(x1-x0)*x+(x0*y1-x1*y0)/(x0-x1)

def f(x, xmax, ymax, plato):

    segments = [x <= plato,
                (plato < x) & (x <= xmax - plato),
                (xmax - plato < x) & (x <= xmax + plato),
                (xmax + plato < x) & (x <= jnp.pi - plato),
                (jnp.pi - plato < x) & (x <= jnp.pi)
                ]

    functions = [line(x, 0, 0, plato, 0),
                 line(x, plato, 0, xmax - plato, ymax),
                 line(x, xmax - plato, ymax, xmax + plato, ymax),
                 line(x, xmax + plato, ymax, jnp.pi - plato, 1),
                 line(x, jnp.pi - plato, 1, jnp.pi, 1)]

    return jnp.piecewise(x, segments, functions)

Now I plot the resulting function using either vmap or a list comprehension to make an array of function values, and the two results do not agree.

xmax = jnp.pi/2
ymax = 2
plato=0.1

g = lambda x: f(x, xmax, ymax, plato)

x_sweep = jnp.linspace(0, jnp.pi, 100)
plt.plot(x_sweep, vmap(g)(x_sweep), label='vmap')  
plt.plot(x_sweep, [g(x) for x in x_sweep], label='list comprehension')
plt.legend()

What could that be?

[soraros]

Quote myself from here:

I think it is a miscompilation bug introduced in #9482, and has been fixed in #9563.