jit[just in time] compilation issue with radius_graph function

[VinaySingh561]

Hi, I am trying to use JIT compilation for the radius graph in my research work. Below is a trivial example demonstrating the problem:
key = jax.random.PRNGKey(0) pos = jax.random.normal(key, (20, 3)) batch = jnp.arange(20) < 10 jax.jit(radius_graph)(pos, 0.8, batch=batch)

And I am getting the following error ::
""""ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: traced array with shape int32[].
The size argument of jnp.nonzero must be statically specified to use jnp.nonzero within JAX transformations.
The error occurred while tracing the function radius_graph at /usr/local/lib/python3.10/dist-packages/e3nn_jax/_src/radius_graph.py:9 for jit. This concrete value was not available in Python because it depends on the values of the arguments pos and r_max.""""

[ameya98]

Thanks for the question. You need to specify a static size so that the output has known shape:

key = jax.random.PRNGKey(0)
pos = jax.random.normal(key, (20, 3))
jax.jit(e3nn.radius_graph, static_argnames=["size"])(pos, r_max=0.8, batch=None, size=10)

but right now it seems like the batch argument is not supported. Let me fix that!
https://colab.research.google.com/drive/1URD6ZX1d2BGIKAK2L59KbzKl7nfwDMI2?usp=sharing

[ameya98]

I would also just use matscipy.neighbours.neighbour_list which is what we do in Symphony: https://github.com/atomicarchitects/symphony/blob/main/symphony/data/input_pipeline.py#L25

[mitkotak]

There's also JAX-MD's optimized neighbor list partitioner if you need to recompute your neighbor list at every step.