[Python] TracerBoolConversion error while attempting to use @jit on functions

In an effort to optimize an existing function used in an optimization algorithm by applying @jit, I encountered some issues. When running the following function:

import jax
import jax.numpy as jnp
from jax import grad, jacobian
from scipy.optimize import minimize
from scipy.interpolate import BSpline

jax.config.update("jax_enable_x64", True)

@jax.jit
def bspline(t_values, knots, coefficients, degree):
"""
Generate a B-spline curve from given knots and coefficients.

Parameters:
- t_values: Array of parameter values where the spline is evaluated.
- knots: Knot vector as a 1D numpy array.
- coefficients: Control points as a 1D numpy array of shape (n,).
- degree: Degree of the B-spline (e.g., 3 for cubic B-splines).

Returns:
- A numpy array of shape (num_points,) representing the B-spline curve.
"""
def basis_function(i, k, t, knots):
"""Compute the basis function recursively."""
if k == 0:
return jnp.where((knots <= t) & (t < knots[i + 1]), 1.0, 0.0)
else:
denom1 = knots[i + k] - knots
denom2 = knots[i + k + 1] - knots[i + 1]

term1 = (t - knots) / denom1 * basis_function(i, k - 1, t, knots) if denom1 != 0 else 0
term2 = (knots[i + k + 1] - t) / denom2 * basis_function(i + 1, k - 1, t, knots) if denom2 != 0 else 0

return term1 + term2

# Compute the B-spline curve points
curve_points = jnp.zeros(len(t_values))
for i in range(len(coefficients)):
v = basis_function(i, degree, t_values, knots)
curve_points = curve_points + v * coefficients

return curve_points


I get the following error:

TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[].

What I've Tried:


After consulting JAX's official documentation about this error (available here), I modified the basis_function to avoid direct boolean checks:

def basis_function(i, k, t, knots):
"""Compute the basis function recursively."""
return jnp.where(
k == 0,
jnp.where((knots <= t) & (t < knots[i + 1]), 1.0, 0.0),
jnp.where(
(knots[i + k] - knots) != 0,
(t - knots) / (knots[i + k] - knots) * basis_function(i, k - 1, t, knots),
0
) +
jnp.where(
(knots[i + k + 1] - knots[i + 1]) != 0,
(knots[i + k + 1] - t) / (knots[i + k + 1] - knots[i + 1]) * basis_function(i + 1, k - 1, t, knots),
0
)
)


However, now I encounter a RecursionError:

RecursionError: maximum recursion depth exceeded in comparison


This recursion issue seems to stem from applying @jit, as it was not present before.

Continue reading...