Questions about control flow and indexed assignment

Hi, I want to implement a finite element solver with jax. I made a quick demo, which runs well without JIT but extremely slow. There are two problems:

1. How do I deal with control flow? For instance, I have if np.isclose(np.abs(np.linalg.det(jacobian), 0.0): raise ValueError(). Since this will run for every element, I cannot just leave it in python. lax.cond seems not suitable for this purpose. Other examples could be calling different material model function, based on element info stored in an array.

2. How do I do indexed assignment? This is a relatively smaller issue. At least I managed doing it with np.where and jax.ops.index_add. But these two methods are creating instead doing in-place update and it requires some helping arrays. lax.dynamic_update_slice can only update a contiguous sub-array. Also it requires creating a sub-array first. Is there something that could do indexed update/add? I’m not very familiar with tensorflow, but it seems you can do something like my_var[4:8].assign(tf.zeros(4)) . This is not exactly what I want, my_var[4:8].add(tf.zeros(4)) would be better, but it looks better at least. Or else, I should not worry about the memory and performance cost of creating a new array?

Thank you!

[shoyer]

Hi, I want to implement a finite element solver with jax

I just wanted to say that I’m very excited about this sort of thing!

In case it helps, we wrote a version simple/naive finite element solver in autograd (which is conceptually quite similar to JAX) in https://github.com/google-research/neural-structural-optimization

[duhd1993]

@joaogui1 Thanks. I haven’t. It will look very ugly. For the first example, it will need nested lax.cond cuz I have many candidate values. Some sorts of lax.switch or dict object would be what I need here.

For the second example, something like tf.debugging.Assert would be more proper than doing it in python. Currently I just commented them out as there’s one element type and no exceptions in my example and shoyer’s paper.