ValueError: Non-hashable static arguments are not supported, as this can lead to unexpected cache-misses.

[mjhoover1]

I am trying to finish up a function, but an issue arose ValueError: Non-hashable static arguments are not supported, as this can lead to unexpected cache-misses. Static argument (index 0) of type <class 'jax.interpreters.partial_eval.DynamicJaxprTracer'> for function inner_loop_body4 is non-hashable. which doesn't seem to align with similar issues usually attempting to dynamically shape an array using a variable input value. Also, note that the function this is wrapped in uses a @partial(jit, static_argnums=(0,)) since the lax.cond depends on step, the prototype of the function is inner_loop_body4(step, i, lane, pos, vel, num_agents_per_road, sum_progress, curr_light_red, road) the problem line is:

carry = lane, leader_dist
lane, leader_dist = lax.cond(step != 0 and step % lane_change_period == 0, body, lambda x: carry, carry)

This will change the lane variable and leader distance it's used for lane changes in a highway simulator, but it never reshapes the lane jnp array. Here is the body:

def body(carry):
            lane, leader_dist = carry
            min_target_lane = lane[road][i] - 1. * sigmoid2(lane[road][i] - 1., .5)
            max_target_lane = lane[road][i] + 1. * (1. - sigmoid2(lane[road][i] - 3 + 1., .5)) # 3=number of lanes
            leader_dists = jnp.zeros(3) # 3=number of lanes
            carry = leader_dists, min_target_lane, max_target_lane
            leader_dists, _, _ = lax.fori_loop(0, 3, inner_loop_body5, carry) # 3=number of lanes
            max_leader_dist = max_pos(leader_dists, 3) # 3=number of lanes
            furthest_ahead_leader_lane = select_x_where_y_equals_z(lane[road], pos[road], pos[road, i] + max_leader_dist, num_agents_per_road)
            new_lane = select_x_if_larger_y_else_z(furthest_ahead_leader_lane, .5, lane[road, i])
            clearance_gain = max_leader_dist - leader_dist
            new_lane_thresh_factor = sigmoid(clearance_gain - min_clearance_gain)
            new_lane_thresh_factor *= sigmoid(vel[road, i] - min_lane_change_vel)
            new_lane = new_lane * new_lane_thresh_factor + lane[road, i] * (1. - new_lane_thresh_factor)
            lane = lane.at[road, i].set(new_lane)
            return lane, leader_dist

This function also relies on inner_loop_body5:

def inner_loop_body5(j, carry):
            leader_dists, min_target_lane, max_target_lane = carry
            leader_dists = leader_dists.at[j].set(min_pos_delta_if_equal(pos[road], num_agents_per_road, i, lane[road], j + 1))
            ir = in_range(j + 1, min_target_lane - .5, max_target_lane + .5)
            leader_dists = leader_dists.at[j].multiply(ir)
            return leader_dists, min_target_lane, max_target_lane

I'm not quite sure what's causing this issue, but is it possible to get around this error? The above is an attempt to summarize the problem to create a complete minimal example. Here is a more verbose example if necessary to get the error to occur, but none of the aux functions are a problem the issues are from the lax.cond call above:
Google colab link: https://colab.research.google.com/drive/1mLIARkjf-ew9WElW-XesQZBqliCTpnzH?usp=sharing
Code:

import jax.numpy as jnp
from jax import jit, lax
from jax.config import config
from jax.scipy.special import logsumexp
import jax
from functools import partial
from jax.scipy.special import expit

config.update('jax_platform_name', 'cpu')
config.update("jax_debug_nans", True)
config.update("jax_enable_x64", True)
config.update('jax_disable_jit', False)

lane_change_period = 25
min_clearance_gain = 10.
min_lane_change_vel = 8.
traffic_sim_sigmoid_slope = 16.0
traffic_light_ahead_sigmoid_slope = 16.0
road_end_pos = 250.
light_pos=100.
road = 0
i = 0
pos = [[40., 300., 340., 380.]]
num_agents_per_road = 4
lane = [[40., 300., 340., 380.]]

@jit
def min_pos_delta_if_equal(x, num_elems, ref_idx, y, z):
    M = 1000.

    @jit
    def body(carry):
        i, sum = carry
        diff = x[i] - x[ref_idx]
        neg_offset = 2 * M * (1.0 - sigmoid(diff))
        diff += neg_offset
        lane_shift = 1000. * (1.0 - in_range(y[i] - z, -0.5, 0.5))
        diff += lane_shift
        exp_arg = -diff
        sum += expit(exp_arg)
        return sum

    @jit
    def outer_body(i, sum):
        return lax.cond(i == ref_idx, lambda x: x[1], body, (i, sum))

    sum = lax.fori_loop(0, num_elems, outer_body, 0.)
    return -jnp.log(sum)

@jit
def sigmoid(x):
    sigmoid_slope = 32.
    return 1.0 / (1.0 + expit(-sigmoid_slope * x))

@jit
def sigmoid2(x, x_offset):
    sigmoid_slope = 32.
    # return my_div(1. + jnp.exp(-sigmoid_slope * (x - x_offset)))
    return 1. / (1. + expit(-sigmoid_slope * (x - x_offset)))

@jit
def in_range_unscaled(x, from_, to):
    return sigmoid2(x, from_)*(1. - sigmoid2(x, to))

@jit
def in_range(x, from_, to):
    return in_range_unscaled(x, from_, to) / in_range_unscaled((from_ + to) / 2.0, from_, to) # second term just for scaling

@jit
def max_pos(x, num_elems):
    return logsumexp(x[:num_elems])

@jit
def select_x_where_y_equals_z(x, y, z, num_elems):
    return (x[:num_elems] * in_range(y[:num_elems] - z, -1., 1.)).sum() # range_scale is 1.

@jit
def select_x_if_larger_y_else_z(x, y, z):
    return sigmoid2(x, y) * x + (1.0 - sigmoid2(x, y)) * z

@jit
def select_x_where_y0_equals_z0_and_y1_equals_z1(x, y0, z0, y1, z1, num_elems):
    return jnp.sum(jnp.where(jnp.arange(len(x)) < num_elems, x, 0) *
    in_range(jnp.where(jnp.arange(len(y0)) < num_elems, y0, 0) - z0, -1., 1.) *
    in_range(jnp.where(jnp.arange(len(y1)) < num_elems, y1, 0) - z1, -.1, .1))

@jit
def sigmoid_with_slope(x, slope):
    return 1.0 / (1.0 + expit(-slope * x))


@partial(jit, static_argnums=(0,))
def inner_loop_body4(step, i, lane, pos, vel, num_agents_per_road, sum_progress, curr_light_red, road):
        def inner_loop_body5(j, carry):
            leader_dists, min_target_lane, max_target_lane = carry
            leader_dists = leader_dists.at[j].set(min_pos_delta_if_equal(pos[road], num_agents_per_road, i, lane[road], j + 1))
            ir = in_range(j + 1, min_target_lane - .5, max_target_lane + .5)
            leader_dists = leader_dists.at[j].multiply(ir)
            return leader_dists, min_target_lane, max_target_lane

        def body(carry):
            lane, leader_dist = carry
            min_target_lane = lane[road][i] - 1. * sigmoid2(lane[road][i] - 1., .5)
            max_target_lane = lane[road][i] + 1. * (1. - sigmoid2(lane[road][i] - 3 + 1., .5)) # 3=number of lanes
            leader_dists = jnp.zeros(3) # 3=number of lanes
            carry = leader_dists, min_target_lane, max_target_lane
            leader_dists, _, _ = lax.fori_loop(0, 3, inner_loop_body5, carry) # 3=number of lanes
            max_leader_dist = max_pos(leader_dists, 3) # 3=number of lanes
            furthest_ahead_leader_lane = select_x_where_y_equals_z(lane[road], pos[road], pos[road, i] + max_leader_dist, num_agents_per_road)
            new_lane = select_x_if_larger_y_else_z(furthest_ahead_leader_lane, .5, lane[road, i])
            clearance_gain = max_leader_dist - leader_dist
            new_lane_thresh_factor = sigmoid(clearance_gain - min_clearance_gain)
            new_lane_thresh_factor *= sigmoid(vel[road, i] - min_lane_change_vel)
            new_lane = new_lane * new_lane_thresh_factor + lane[road, i] * (1. - new_lane_thresh_factor)
            lane = lane.at[road, i].set(new_lane)
            return lane, leader_dist

        leader_dist = min_pos_delta_if_equal(pos[road], num_agents_per_road, i, lane[road], lane[road][i])
        carry = lane, leader_dist
        lane, leader_dist = lax.cond(step != 0 and step % lane_change_period == 0, body, lambda x: carry, carry)
        
        leader_dist = min_pos_delta_if_equal(pos[road], num_agents_per_road, i, lane[road], lane[road, i]) # updated after lane change
        leader_vel = select_x_where_y0_equals_z0_and_y1_equals_z1(vel[road], lane[road], lane[road][i], pos[road], pos[road][i] + leader_dist, num_agents_per_road)
        light_dist = light_pos - pos[road, i]
        light_ahead = sigmoid_with_slope(light_dist, 16.) # traffic_light_ahead_sigmoid_slope=16.
        braking_at_light = curr_light_red * light_ahead * sigmoid_with_slope(leader_dist - light_dist, 16.) # traffic_sim_sigmoid_slope=16.
        leader_dist = braking_at_light * light_dist + (1. - braking_at_light) * leader_dist
        leader_vel = (1. - braking_at_light) * leader_vel
        accel = idm(vel[road, i], leader_dist, leader_vel, 13.888, 2.) # 13.888=speed limit, max_accel=2.
        vel = vel.at[road, i].add(accel * .1) # .1=tau=time step
        vel = vel.at[road, i].multiply(sigmoid_with_slope(vel[road, i], traffic_sim_sigmoid_slope)) # allow velocities > 0 - eps only
        progress = vel[road, i] * .1 # .1=tau=time step
        pos = pos.at[road, i].add(progress)
        sum_progress += progress
        past_road_end = sigmoid_with_slope(pos[road,i] - road_end_pos, traffic_sim_sigmoid_slope)
        pos = pos.at[road, i].add(-pos[road,i] * past_road_end)
        return lane, pos, vel, num_agents_per_road, sum_progress, curr_light_red, road

@jit
def wrapper_inner_loop_body4(i, carry):
        step, lane, pos, vel, num_agents_per_road, sum_progress, curr_light_red, road = carry
        lane, pos, vel, num_agents_per_road, sum_progress, curr_light_red, road = inner_loop_body4(step, i, lane, pos, vel, num_agents_per_road, sum_progress, curr_light_red, road)
        return step, lane, pos, vel, num_agents_per_road, sum_progress, curr_light_red, road


lax.fori_loop(0, num_agents_per_road - 3, wrapper_inner_loop_body4, (1, lane, pos, [[13.888, 13.888, 13.888, 13.888]], num_agents_per_road, 0.0, 0.63999973, road)) # 3 is number of lanes

The first issue I was getting was The problem arose with the `bool` function. While tracing the function while_body_fun at which makes sense since lax.cond depends on the step and i, so I made them static using @partial(jit, static_argnums=(0,1)). Now, I'm getting this issue which I can't figure out how to get around, thanks!

[jakevdp]

You jit-compile wrapper_inner_loop_body4, which means all arguments passed to it are non-static, including the first element of carry, which is step.

You then pass step as the first argument to inner_loop_body4, which is marked as static, and this leads to an error. To fix that, you should avoid marking this argument as static. The reason this causes issues is because you're using Python's and operator with traced variables; and attempts to eagerly convert outputs to static booleans (a behavior that is impossible to override). For that reason, JAX follows NumPy and uses the elemenwise and operator (&) for this type of operation:

leader_dist = min_pos_delta_if_equal(pos[road], num_agents_per_road, i, lane[road], lane[road][i])

At this point you get a new error because you're attempting to index a list with a traced value – you can fix that by casting the lists to a JAX array:

pos = jnp.array([[40., 300., 340., 380.]])
lane = jnp.array([[40., 300., 340., 380.]])

Then you hit an IndexError because you're attempting to create a dynamically-sized array within a traced function. JAX arrays must be statically shaped, so you'll need to find a different way to express this logic if you want to use it with JIT.

Side-note: there's something strange in your sigmoid functions. expit(x) is equivalent to sigmoid(x), so 1 / (1 + expit(-x)) is something else entirely.

Best of luck!

[mjhoover1]

First and foremost, I want to thank you greatly @jakevdp. I can't tell you how long I've been working on this simulator to get it to be differentiable as well as jittable, and couldn't have done it without your help. To make this post most informative for others, I'll clarify a few things. I knew about wanting to use & instead of and for JAX but it was computing the conditional statement differently than using and the reason is I was not adding parenthesis doing this step != 0 & step % lane_change_period == 0 instead of this (step != 0) & (step % lane_change_period == 0) so make sure to add parenthesis. As @jakevdp noted above I removed the @partial(jit, static_argnums=0) above my inner_loop_body4 for exactly the reason mentioned above. I changed my lax.cond(step != 0 and step % lane_change_period == 0, body, lambda x: carry, carry) to lax.cond((step != 0) & (step % lane_change_period == 0), body, lambda x: carry, carry). This is when the IndexError occurs which was due to splicing a JAX array that must be statically shaped, so in times when you are using x[:number]. The best approach I recommend is that you or even JAX make some type of splice function which does about the same thing while obeying JAX rules such as:

@jit
def splice(x, number): # splice is almost equivalent to x[:number] but doesn't reshape JAX arrays to attempt dynamically sizing
    return jnp.where(jnp.arange(len(x)) < number, x, 0)
import numpy as np
q = np.array([1, 2, 3, 4])
print(q[:2]) # output: array([1, 2])
print(splice(q, 2) # output: DeviceArray([1, 2, 0, 0], dtype=int64)

This was the issue that was causing the IndexError specifically it was this function:

def max_pos(x, num_elems):
    return logsumexp(x[:num_elems])

As well as a couple of other functions that I was using Python splice on a JAX array, to fix this I made a splice function as described above and changed this function to:

@jit
def max_pos(x, num_elems):
    return logsumexp(splice(x, num_elems))

Then voila the code worked as expected, as for the sigmoid and expit(x) situation I thought the expit(x) was equivalent to a jnp.exp(x). However, as @jakevdp explicitly stated this is not the case and that's what happens after coding for 8 hours I suppose. For future reference, I will link to the Google Colab again which will contain the original code that didn't work as well as the fixed code, as well as post the fixed code here too. I hope this helps anyone else who comes across these errors, thanks!

Edit: A note was mentioned below by @jakevdp which is a valid point which I mentioned how to update the default value for different usage please check the comment below for details but the updated code will be updated to satisfy this point. So, that the splice function is more versatile for different cases.

Colab link: https://colab.research.google.com/drive/1mLIARkjf-ew9WElW-XesQZBqliCTpnzH?usp=sharing
Updated code:

import jax.numpy as jnp
from jax import jit, lax
# again, this only works on startup!
from jax.config import config
from jax.scipy.special import logsumexp
import jax
from functools import partial
from jax.scipy.special import expit

# import autograd.numpy as jnp   # Thinly-wrapped version of Numpy
# from autograd import grad

debug = True

config.update('jax_platform_name', 'cpu')
config.update("jax_debug_nans", True)
config.update("jax_enable_x64", True)
config.update('jax_disable_jit', False)

lane_change_period = 25
min_clearance_gain = 10.
min_lane_change_vel = 8.
traffic_sim_sigmoid_slope = 16.0
traffic_light_ahead_sigmoid_slope = 16.0
road_end_pos = 250.
light_pos=100.
tau = .1
light_red_duration = 5.

@jit
def sigmoid(x):
    sigmoid_slope = 32.
    return expit(sigmoid_slope * x)

@jit
def sigmoid2(x, x_offset):
    sigmoid_slope = 32.
    return expit(sigmoid_slope * (x-x_offset))

@jit
def in_range_unscaled(x, from_, to):
    return sigmoid2(x, from_)*(1. - sigmoid2(x, to))

@jit
def in_range(x, from_, to):
    return in_range_unscaled(x, from_, to) / in_range_unscaled((from_ + to) / 2.0, from_, to) # second term just for scaling

@jit
def sigmoid_with_slope(x, slope):
    return expit(slope * x)

@jit
def periodic_step(x, period, offset, slope):
    return sigmoid_with_slope(jnp.sin((x - offset) * jnp.pi / period), slope)

def light_red(x, on_period, offset, slope = -1.0):
    if slope < 0.0:
      slope = 32.
    return periodic_step(x, on_period, offset, slope)

# For differentiable vehicle sorting:
# Given an unordered list of vehicles, we determine a vehicle’s 
# leader by iterating across all vehicles and determining the 
# vehicle with the minimum positive position delta. Since we 
# require both the position and velocity of the leader, we arrive 
# at a two-step process: first, we determine the leader’s position 
# by repeatedly applying smooth minimum as described in Section 3.3, 
# masking negative position deltas using smooth threshold. Finally, 
# select by attribute determines the leader’s velocity based on its 
# position and lane.
@jit
def min_pos_delta_if_equal(x, num_elems, ref_idx, y, z):
    M = 1000.

    @jit
    def body(carry):
        i, sum = carry
        diff = x[i] - x[ref_idx]
        neg_offset = 2 * M * (1.0 - sigmoid(diff))
        diff += neg_offset
        lane_shift = 1000. * (1.0 - in_range(y[i] - z, -0.5, 0.5))
        diff += lane_shift
        exp_arg = -diff
        sum += jnp.exp(exp_arg)
        return sum

    @jit
    def outer_body(i, sum):
        return lax.cond(i == ref_idx, lambda x: x[1], body, (i, sum))

    sum = lax.fori_loop(0, num_elems, outer_body, 0.)
    return -jnp.log(sum)

@jit
def splice(x, num_elems, default_value): # splice is equivalent to x[:num_elems] but works for JAX arrays to dynamically size
    return jnp.where(jnp.arange(len(x)) < num_elems, x, default_value)

@jit
def max_pos(x, num_elems):
    return logsumexp(splice(x, num_elems, -jnp.inf))

@jit
def select_x_where_y_equals_z(x, y, z, num_elems):
    return (splice(x, num_elems, 0) * in_range(splice(y, num_elems, 0) - z, -1., 1.)).sum() # range_scale is 1.

# @functools.partial(jax.jit, static_argnums=(5,))
@jit
def select_x_where_y0_equals_z0_and_y1_equals_z1(x, y0, z0, y1, z1, num_elems):
    return jnp.sum(splice(x, num_elems, 0) *
    in_range(splice(y0, num_elems, 0) - z0, -1., 1.) *
    in_range(splice(y1, num_elems, 0) - z1, -.1, .1))

@jit
def select_x_if_larger_y_else_z(x, y, z):
    return sigmoid2(x, y) * x + (1.0 - sigmoid2(x, y)) * z

@jit
def idm(vel, leader_dist, leader_vel, speed_limit, max_accel):
    s0 = 5.0  # min spacing
    t = 3.0;   # desired time headway
    accel_term0 = jnp.power(vel / speed_limit, 4)
    accel_term1 = jnp.power((s0 + vel * t + (vel * (vel - leader_vel) / (2 * max_accel))) / leader_dist, 2)
    return max_accel * (1.0 - accel_term0 - accel_term1)


def traffic_sim_smoothed_jax(x, num_agents_per_road):
    num_agents_per_road += 3 # 3 is number of lanes
    pos = jnp.zeros((1, num_agents_per_road)) # 1 is number of roads
    vel = jnp.zeros((1, num_agents_per_road)) # 1 is number of roads
    lane = jnp.zeros((1, num_agents_per_road)) # 1 is number of roads

    @jit
    def inner_loop_body1(carry, i):
        lane, pos, vel, road = carry
        idx = num_agents_per_road - 3 + i - 1 # 3=number of lanes
        lane = lane.at[road, idx].set(1 + (i - 1) % 3) # 3=number of lanes
        pos = pos.at[road, idx].set(300. + (i - 1) * 40) # 300. is dummy car pos, & 40 is init_pos_inc
        vel = vel.at[road, idx].set(13.888) # 13.888=speed limit
        return (lane, pos, vel, road), None

    @jit
    def inner_loop_body2(carry, i):
        lane, pos, vel, road = carry
        lane = lane.at[road, i].set(1. + i % 3) # 3 is number of 
        pos = pos.at[road, i].set(((i + 1) * 40) % int(250.)) # 40=init_pos_inc, 250.=
        vel = vel.at[road, i].set(13.888) # 13.888 is the speed limit
        return (lane, pos, vel, road), None

    @jit
    def outer_loop_body1(carry, road):
        lane, pos, vel = carry
        carry, _ = lax.scan(inner_loop_body1, (lane, pos, vel, road), jnp.arange(1, 3 + 1)) # 3 = number of lanes need +1 for while i <= 3
        # carry = lax.fori_loop(1, 3 + 1, inner_loop_body1, (lane, pos, vel, road)) # 3 = number of lanes need +1 for while i <= 3
        carry, _ = lax.scan(inner_loop_body2, carry, jnp.arange(num_agents_per_road - 3)) # 3=number of lanes
        # carry = lax.fori_loop(0, num_agents_per_road - 3, inner_loop_body2, carry) # 3=number of lanes
        lane, pos, vel, road = carry
        return (lane, pos, vel), None


    res, _ = lax.scan(outer_loop_body1, (lane, pos, vel), jnp.arange(1 + 1)) # 1 is the num_roads need + 1 since not zero indexed
    lane, pos, vel = res

    @jit
    def inner_loop_body4(step, i, lane, pos, vel, num_agents_per_road, sum_progress, curr_light_red, road):
        def inner_loop_body5(j, carry):
            leader_dists, min_target_lane, max_target_lane = carry
            leader_dists = leader_dists.at[j].set(min_pos_delta_if_equal(pos[road], num_agents_per_road, i, lane[road], j + 1))
            ir = in_range(j + 1, min_target_lane - .5, max_target_lane + .5)
            leader_dists = leader_dists.at[j].multiply(ir)
            return leader_dists, min_target_lane, max_target_lane

        def body(carry):
            lane, leader_dist = carry
            min_target_lane = lane[road][i] - 1. * sigmoid2(lane[road][i] - 1., .5)
            max_target_lane = lane[road][i] + 1. * (1. - sigmoid2(lane[road][i] - 3 + 1., .5)) # 3=number of lanes
            leader_dists = jnp.zeros(3) # 3=number of lanes
            carry = leader_dists, min_target_lane, max_target_lane
            leader_dists, _, _ = lax.fori_loop(0, 3, inner_loop_body5, carry) # 3=number of lanes
            max_leader_dist = max_pos(leader_dists, 3) # 3=number of lanes
            furthest_ahead_leader_lane = select_x_where_y_equals_z(lane[road], pos[road], pos[road, i] + max_leader_dist, num_agents_per_road)
            new_lane = select_x_if_larger_y_else_z(furthest_ahead_leader_lane, .5, lane[road, i])
            clearance_gain = max_leader_dist - leader_dist
            new_lane_thresh_factor = sigmoid(clearance_gain - min_clearance_gain)
            new_lane_thresh_factor *= sigmoid(vel[road, i] - min_lane_change_vel)
            new_lane = new_lane * new_lane_thresh_factor + lane[road, i] * (1. - new_lane_thresh_factor)
            lane = lane.at[road, i].set(new_lane)
            return lane, leader_dist

        leader_dist = min_pos_delta_if_equal(pos[road], num_agents_per_road, i, lane[road], lane[road][i])
        carry = lane, leader_dist
        # lane, leader_dist = lax.cond(step != 0 and step % lane_change_period == 0, body, lambda x: carry, carry)
        lane, leader_dist = lax.cond((step != 0) & (step % lane_change_period == 0), body, lambda x: carry, carry)
        
        leader_dist = min_pos_delta_if_equal(pos[road], num_agents_per_road, i, lane[road], lane[road, i]) # updated after lane change
        leader_vel = select_x_where_y0_equals_z0_and_y1_equals_z1(vel[road], lane[road], lane[road][i], pos[road], pos[road][i] + leader_dist, num_agents_per_road)
        light_dist = light_pos - pos[road, i]
        light_ahead = sigmoid_with_slope(light_dist, traffic_light_ahead_sigmoid_slope) 
        braking_at_light = curr_light_red * light_ahead * sigmoid_with_slope(leader_dist - light_dist, traffic_sim_sigmoid_slope) 
        leader_dist = braking_at_light * light_dist + (1. - braking_at_light) * leader_dist
        leader_vel = (1. - braking_at_light) * leader_vel
        accel = idm(vel[road, i], leader_dist, leader_vel, 13.888, 2.) # 13.888=speed limit, max_accel=2.
        vel = vel.at[road, i].add(accel * tau) 
        vel = vel.at[road, i].multiply(sigmoid_with_slope(vel[road, i], traffic_sim_sigmoid_slope)) # allow velocities > 0 - eps only
        progress = vel[road, i] * tau 
        pos = pos.at[road, i].add(progress)
        sum_progress += progress
        past_road_end = sigmoid_with_slope(pos[road,i] - road_end_pos, traffic_sim_sigmoid_slope)
        pos = pos.at[road, i].add(-pos[road,i] * past_road_end)
        return lane, pos, vel, num_agents_per_road, sum_progress, curr_light_red, road

    @jit
    def wrapper_inner_loop_body4(i, carry):
        step, lane, pos, vel, num_agents_per_road, sum_progress, curr_light_red, road = carry
        lane, pos, vel, num_agents_per_road, sum_progress, curr_light_red, road = inner_loop_body4(step, i, lane, pos, vel, num_agents_per_road, sum_progress, curr_light_red, road)
        return step, lane, pos, vel, num_agents_per_road, sum_progress, curr_light_red, road

    @jit
    def inner_loop_body3(road, carry):
        step, lane, pos, vel, num_agents_per_road, sum_progress, curr_light_red = carry
        step, lane, pos, vel, num_agents_per_road, sum_progress, curr_light_red, road = lax.fori_loop(0, 
        num_agents_per_road - 3, wrapper_inner_loop_body4, (step, lane, pos, vel, num_agents_per_road, sum_progress, curr_light_red, road)) # 3 is number of lanes
        return step, lane, pos, vel, num_agents_per_road, sum_progress, curr_light_red

    @jit
    def outer_loop_body2(carry):
        step, lane, pos, vel, num_agents_per_road, sum_progress = carry
        curr_light_red = light_red(tau * step, light_red_duration, x, traffic_sim_sigmoid_slope) # x=light_offset
        carry = step, lane, pos, vel, num_agents_per_road, sum_progress, curr_light_red
        step, lane, pos, vel, num_agents_per_road, sum_progress, _ = lax.fori_loop(0, 1, inner_loop_body3, carry) # 1 is number of roads
        return step + 1, lane, pos, vel, num_agents_per_road, sum_progress

    cond_fun = lambda carry: .1 * carry[0] <= 10.
    _, _, _, _, _, sum_progress = lax.while_loop(cond_fun, outer_loop_body2, (0, lane, pos, vel, num_agents_per_road, 0.0))
    return sum_progress


if __name__ == "__main__":
    num_agents_per_road = 1
    tau = .1
    x_init = tau/4 # Initial x time step
    print("x, y, dydx")
    value = traffic_sim_smoothed_jax(x_init, num_agents_per_road)
    gradient = jax.jacfwd(lambda x: traffic_sim_smoothed_jax(x, num_agents_per_road))(x_init)
    # value, gradient = value_and_grad(traffic_sim_smoothed_jax)(x_init, num_agents_per_road)
    print(str(x_init) + ", ", value, ", " + str(gradient))
    # print(grad(traffic_sim_smoothed_autograd)(x_init, num_agents_per_road))

[jakevdp]

One other note: I don't think logsumexp(x[:N]) is equivalent to logsumexp(splice(x, N)). I believe you need the splice to default to -np.inf rather than 0 for that to work.

[mjhoover1]

Good catch, it does seem to be outputting the correct value. Although, I trust your intuition in order to satisfy this requirement here is what I did for future reference:

@jit
def splice(x, num_elems, default): # splice is equivalent to x[:num_elems] but works for JAX arrays to dynamically size
    return jnp.where(jnp.arange(len(x)) < num_elems, x, default)

Now, for each specific case the default value can be changed as done here:

@jit
def max_pos(x, num_elems):
    return logsumexp(splice(x, num_elems, -jnp.inf))

@jit
def select_x_where_y_equals_z(x, y, z, num_elems):
    return (splice(x, num_elems, 0) * in_range(splice(y, num_elems, 0) - z, -1., 1.)).sum() # range_scale is 1.

@jit
def select_x_where_y0_equals_z0_and_y1_equals_z1(x, y0, z0, y1, z1, num_elems):
    return jnp.sum(splice(x, num_elems, 0) *
    in_range(splice(y0, num_elems, 0) - z0, -1., 1.) *
    in_range(splice(y1, num_elems, 0) - z1, -.1, .1))

I will make these updates in the code above as well as the Colab notebook, thank you for your time!