def test_fenics_vjp():
numpy_output, fenics_output, fenics_inputs, tape = evaluate_primal(
solve_fenics, templates, *inputs)
g = np.ones_like(numpy_output)
vjp_out = evaluate_pullback(fenics_output, fenics_inputs, tape, g)
check1 = np.isclose(vjp_out[0], np.asarray(-2.91792642))
check2 = np.isclose(vjp_out[1], np.asarray(2.43160535))
assert check1 and check2
def test_firedrake_vjp():
numpy_output, firedrake_output, firedrake_inputs, tape = evaluate_primal(
solve_firedrake, templates, *inputs)
g = np.ones_like(numpy_output)
vjp_out = evaluate_pullback(firedrake_output, firedrake_inputs, tape, g)
check1 = np.isclose(vjp_out[0], np.asarray(-1.13533304))
check2 = np.isclose(vjp_out[1], np.asarray(0.94611087))
assert check1 and check2
def perform(self, node, inputs, outputs, params):
Δfenics_output = inputs[0]
fenics_output, fenics_inputs, tape = params
numpy_grads = evaluate_pullback(fenics_output, fenics_inputs, tape,
Δfenics_output)
theano_grads = (theano.gradient.grad_undefined(self, i, inputs[i])
if ng is None else ng
for i, ng in enumerate(numpy_grads))
for i, tg in enumerate(theano_grads):
outputs[i][0] = tg
def vjp_fun1_p_impl(aux_args, g):
fe_aux, args = aux_args
fenics_output, fenics_inputs, tape = (
fe_aux.fenics_output,
fe_aux.fenics_inputs,
fe_aux.tape,
)
return tuple(
vjp if vjp is not None else jax.ad_util.zeros_like_jaxval(args[i])
for i, vjp in enumerate(
evaluate_pullback(fenics_output, fenics_inputs, tape, g)
)
)
def test_vjp_assemble_eval():
numpy_output, fenics_output, fenics_inputs, tape = evaluate_primal(
assemble_fenics, templates, *inputs)
g = np.ones_like(numpy_output)
vjp_out = evaluate_pullback(fenics_output, fenics_inputs, tape, g)
fdm_jac0 = fdm.jacobian(ff0)(inputs[0])
fdm_jac1 = fdm.jacobian(ff1)(inputs[1])
fdm_jac2 = fdm.jacobian(ff2)(inputs[2])
check1 = np.allclose(vjp_out[0], fdm_jac0)
check2 = np.allclose(vjp_out[1], fdm_jac1)
check3 = np.allclose(vjp_out[2], fdm_jac2)
assert check1 and check2 and check3
def test_theano_vjp():
theano.config.compute_test_value = "ignore"
numpy_output, fenics_output, fenics_inputs, tape = evaluate_primal(
assemble_firedrake, templates, *inputs)
vjp_op = FenicsVJPOp(assemble_firedrake, templates, fenics_output,
tuple(fenics_inputs), tape)
g = theano.tensor.vector()
f = theano.function([g], vjp_op(g))
theano_output = f(np.ones(1))
numpy_output = evaluate_pullback(fenics_output, tuple(fenics_inputs), tape,
np.ones(1))
for to, no in zip(theano_output, numpy_output):
with check:
assert np.allclose(to, no)