Haoliang Wang

vmap for `jax_kernel` doesn't seem to work either: I'm taking [this example from the documentation](https://nvidia.github.io/warp/modules/interoperability.html#id1), and tried to vmap it to a 2d array ``` import jax import jax.numpy as...

Just wanted to kindly follow up on this—do you have an estimate of when the bug might be fixed? Really appreciate your help! @nvlukasz @shi-eric