vmap for build_jax_fem_eval?

[nguyenvanhaibk92]

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.