Geometry Optimization

[andrinr]

I am currently trying to perform a geometry optimization using JAX-FEM. However when I set the mesh vertices to a jax array, which is necessary to get gradients for the vertex positions, I get the error:

  File "/home/andrin/miniconda3/envs/fem/lib/python3.12/site-packages/jax/_src/numpy/array_methods.py", line 1067, in meth
    return getattr(self.aval, name).fun(self, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/andrin/miniconda3/envs/fem/lib/python3.12/site-packages/jax/_src/numpy/array_methods.py", line 384, in _take
    return lax_numpy.take(self, indices, axis=axis, out=out, mode=mode, unique_indices=unique_indices,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/andrin/miniconda3/envs/fem/lib/python3.12/site-packages/jax/_src/numpy/lax_numpy.py", line 11468, in take
    return _take(a, indices, None if axis is None else operator.index(axis), out,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/andrin/miniconda3/envs/fem/lib/python3.12/site-packages/jax/_src/numpy/lax_numpy.py", line 11494, in _take
    raise NotImplementedError("The 'raise' mode to jnp.take is not supported.")
NotImplementedError: The 'raise' mode to jnp.take is not supported.