ivanyashchuk / jax-fenics-adjoint Goto Github PK

Something has changed in JAX and now tests fail because the returned jvp tangent contains nans in place of zeros.

Examples use old function names and should be updated.

It's a good idea to add to CI with a maximum number of iterations to test they run without fails.

Hi Ivan,
in the notebook, we have the function

@build_jax_fem_eval(solve_templates)
def fenics_solve(f):

I try to use jax.vmap. It works in the test files, where we just want to compute outputs of u. However, inside the neural network

def eval_nn(net_params):
f_nn = np.ravel(net_apply(net_params, W.tabulate_dof_coordinates()))
u = vmap(fenics_solve(f_nn))
norm_u = np.linalg.norm(u)
return norm_u

but it does not work. For example, in the case that we deal with multiple samples.
Can you give me further guidance?

Thank you.

This commit google/jax@632876d removed defvjp_all and defjvp_all functions from JAX.

I need to update the way we register custom derivatives.

Hi @IvanYashchuk . Thank you for creating this amazing interface for JAX with FEniCS !! I am receiving the following error and was hoping you could help me out with it.

Error occurs when I run my code with @jit

Please let me know if you require anymore information.