Repository navigation
Fix type casting for auto-differentiability #1267
Description
Add your issue here
Description
During #1251 it was discovered that converting JAX arrays to base types though function such as float breaks auto-differentiability. We should instead use the Array API compatible functions:
astype: For examplefloat(x)would becomexp.astype(x, float).isnan: For examplemath.isnan(x)woudl becomexp.isnan(x).
Example error
> ws.append(RadialWindow(z, w, float(zmid)))
^^^^^^^^^^^
E jax.errors.ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: traced array with shape float64[]
E The problem arose with the `float` function. If trying to convert the data type of a value, try using `x.astype(float)` or `jnp.array(x, float)` instead.
E
E See https://docs.jax.dev/en/latest/errors.html#jax.errors.ConcretizationTypeError
E --------------------
E For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.
glass/shells.py:420: ConcretizationTypeError
Activity