initial implementation

fem/dfem/doperator.hpp

@@ -31,6 +31,10 @@ namespace mfem::future
 using action_t =
    std::function<void(std::vector<Vector> &, const std::vector<Vector> &, Vector &)>;
 
+/// @brief Type alias for a function that computes the VJP
+using vjp_action_t =
+   std::function<void(std::vector<Vector> &, const std::vector<Vector> &, const Vector &, std::vector<Vector> &)>;
+
 /// @brief Type alias for a function that computes the cache for the action of a derivative
 using derivative_setup_t =
    std::function<void(std::vector<Vector> &, const Vector &)>;
@@ -280,6 +284,68 @@ public:
       mult_level = level;
    }
 
+   /// @brief Compute the vector-Jacobian product using reverse mode differentiation in Enzyme.
+   ///
+   /// @param solutions_in The solution vector in which to compute the action.
+   /// @param v_in The vector to multiply with the transposed Jacobian.
+   /// @param vjp_out Result vector of the vector-Jacobian product.
+   void vjp(const Vector &solutions_in, const Vector &v_in, Vector &vjp_out) const
+   {
+      MFEM_ASSERT(!vjp_callbacks.empty(), "no integrators have been set");
+
+      if (mult_level == MultLevel::LVECTOR)
+      {
+         get_lvectors(solutions, solutions_in, solutions_l);
+         vjp_out = 0.0;
+         for (auto &vjp_cb : vjp_callbacks)
+         {
+            vjp_cb(solutions_l, parameters_l, v_in, vjp_l);
+         }
+         // Assuming vjp_out is L-vector for LVECTOR mult_level, but vjp_l is std::vector<Vector>.
+         // Actually, if mult_level == LVECTOR, vjp_out should be a single block L-vector?
+         // In typical usage, LVECTOR means the inputs and outputs are L-vectors.
+         MFEM_ABORT("vjp with LVECTOR mult_level is not fully implemented");
+      }
+      else
+      {
+         prolongation(solutions, solutions_in, solutions_l);
+
+         if (vjp_l.size() != solutions_l.size())
+         {
+            vjp_l.resize(solutions_l.size());
+            for (size_t i = 0; i < solutions_l.size(); i++)
+            {
+               vjp_l[i].SetSize(solutions_l[i].Size());
+            }
+         }
+
+         for (size_t i = 0; i < vjp_l.size(); i++)
+         {
+            vjp_l[i] = 0.0;
+         }
+
+         Vector v_l;
+         prolongation(fields[test_space_field_idx], v_in, v_l);
+
+         for (auto &vjp_cb : vjp_callbacks)
+         {
+            vjp_cb(solutions_l, parameters_l, v_l, vjp_l);
+         }
+
+         // Combine vjp_l back to vjp_out
+         vjp_out = 0.0;
+         int offset = 0;
+         for (size_t i = 0; i < solutions.size(); i++)
+         {
+            const auto P = get_prolongation(solutions[i]);
+            const int width = P->Width();
+            Vector vjp_out_i(vjp_out, offset, width);
+            P->MultTranspose(vjp_l[i], vjp_out_i);
+            offset += width;
+         }
+      }
+   }
+
    /// @brief Compute the action of the operator on a given vector.
    ///
    /// @param solutions_in The solution vector in which to compute the action.
@@ -482,6 +548,7 @@ private:
    MultLevel mult_level = TVECTOR;
 
    std::vector<action_t> action_callbacks;
+   std::vector<vjp_action_t> vjp_callbacks;
    std::map<size_t, std::vector<derivative_setup_t>> derivative_setup_callbacks;
    std::map<size_t,
        std::vector<derivative_action_t>> derivative_action_callbacks;
@@ -502,6 +569,7 @@ private:
    mutable std::vector<Vector> solutions_l;
    mutable std::vector<Vector> parameters_l;
    mutable Vector residual_l;
+   mutable std::vector<Vector> vjp_l;
 
    mutable std::vector<Vector> fields_e;
    mutable Vector residual_e;
@@ -881,6 +949,130 @@ void DifferentiableOperator::AddIntegrator(
       output_restriction_transpose(residual_e, res);
    });
 
+   auto vjp_shmem_info =
+      get_shmem_info<entity_t, num_fields, num_inputs, num_outputs>
+      (input_dtq_maps, output_dtq_maps, fields, num_entities, inputs, num_qp,
+       input_size_on_qp, residual_size_on_qp, element_dof_ordering, 0);
+
+   std::vector<int> sol_field_sizes(solutions.size());
+   for (size_t i = 0; i < solutions.size(); ++i) {
+       sol_field_sizes[i] = get_restriction<entity_t>(solutions[i], element_dof_ordering)->Height() / num_entities;
+   }
+
+   vjp_callbacks.push_back(
+      [
+         dimension, num_entities, num_test_dof, num_qp, q1d, residual_size_on_qp,
+         test_vdim, test_op_dim, inputs, attributes, ir_weights, use_sum_factorization,
+         input_dtq_maps, output_dtq_maps, input_to_field, output_fop, qfunc,
+         thread_blocks, shmem_info = vjp_shmem_info, elem_attributes, element_dof_ordering,
+         output_e_size = output_e_size, test_space_field = fields[test_space_field_idx],
+         num_solutions = solutions.size(),
+         sol_fields = solutions,
+         sol_field_sizes, outputs,
+         &restriction_cb = this->restriction_callback,
+         &fields_e = this->fields_e
+      ]
+      (std::vector<Vector> &sol, const std::vector<Vector> &par, const Vector &v_l, std::vector<Vector> &vjp_l) mutable
+   {
+#ifdef MFEM_USE_ENZYME
+      restriction_cb(sol, par, fields_e);
+
+      Vector v_e(output_e_size);
+      restriction<entity_t>(test_space_field, v_l, v_e, element_dof_ordering);
+
+      std::vector<Vector> vjp_e(num_solutions);
+      for (size_t i = 0; i < num_solutions; i++) {
+         vjp_e[i].SetSize(sol_field_sizes[i] * num_entities);
+         vjp_e[i] = 0.0;
+      }
+
+      auto wrapped_fields_e = wrap_fields(fields_e, shmem_info.field_sizes, num_entities);
+      auto wrapped_v_e = Reshape(v_e.ReadWrite(), output_e_size / num_entities, num_entities);
+
+      const bool has_attr = attributes.Size() > 0;
+      const auto d_attr = attributes.Read();
+      const auto d_elem_attr = elem_attributes->Read();
+
+      Vector shmem_cache(shmem_info.total_size);
+
+      std::array<DeviceTensor<2>, 5> wrapped_vjp_e;
+      for (size_t i = 0; i < num_solutions; i++) {
+         wrapped_vjp_e[i] = Reshape(vjp_e[i].ReadWrite(), sol_field_sizes[i], num_entities);
+      }
+
+      forall([=] MFEM_HOST_DEVICE (int e, void *shmem)
+      {
+         if (has_attr && !d_attr[d_elem_attr[e] - 1]) { return; }
+
+         DeviceTensor<2> dummy_dir_e;
+         auto [input_dtq_shmem_, output_dtq_shmem_, fields_shmem_, direction_shmem_, input_shmem_,
+               shadow_shmem_, residual_shmem_, scratch_shmem_] =
+                  unpack_shmem(shmem, shmem_info, input_dtq_maps, output_dtq_maps,
+                               wrapped_fields_e, dummy_dir_e, num_qp, e);
+         auto &input_dtq_shmem = input_dtq_shmem_;
+         auto &output_dtq_shmem = output_dtq_shmem_;
+         auto &fields_shmem = fields_shmem_;
+         auto &input_shmem = input_shmem_;
+         auto &shadow_shmem = shadow_shmem_;
+         auto &residual_shmem = residual_shmem_;
+         auto &scratch_shmem = scratch_shmem_;
+
+         map_fields_to_quadrature_data(
+            input_shmem, fields_shmem, input_dtq_shmem, input_to_field, inputs, ir_weights,
+            scratch_shmem, dimension, use_sum_factorization);
+
+         set_zero(shadow_shmem);
+
+         auto v_e_i = Reshape(&wrapped_v_e(0, e), output_e_size / num_entities);
+         
+         std::array<DeviceTensor<2>, 1> res_qp_arr = { residual_shmem };
+         std::array<DeviceTensor<1>, 1> v_e_arr = { v_e_i };
+         std::array<size_t, 1> dummy_map = { 0 };
+
+         map_fields_to_quadrature_data(
+            res_qp_arr, v_e_arr, output_dtq_shmem, dummy_map, outputs, ir_weights,
+            scratch_shmem, dimension, use_sum_factorization);
+
+         mfem::future::call_qfunction_vjp<qf_param_ts>(
+            qfunc, input_shmem, shadow_shmem, residual_shmem,
+            residual_size_on_qp, num_qp, q1d, dimension, use_sum_factorization);
+
+         for_constexpr<num_inputs>([&](auto s)
+         {
+            if (input_to_field[s] < num_solutions)
+            {
+               const int trial_vdim = get<s>(inputs).vdim;
+               const int trial_op_dim = get<s>(inputs).size_on_qp / trial_vdim;
+               auto d_qp = Reshape(&(shadow_shmem[s])[0], trial_vdim, trial_op_dim, num_qp);
+               const int num_trial_dof = sol_field_sizes[input_to_field[s]] / trial_vdim;
+               
+               auto y_e_2d = Reshape(&(input_shmem[s])(0, 0), num_trial_dof, trial_vdim);
+
+               map_quadrature_data_to_fields(
+                  y_e_2d, d_qp, get<s>(inputs), input_dtq_shmem[s],
+                  scratch_shmem, dimension, use_sum_factorization);
+                  
+               auto dest = Reshape(&wrapped_vjp_e[input_to_field[s]](0, e), sol_field_sizes[input_to_field[s]]);
+               for(int vd=0; vd<trial_vdim; ++vd) {
+                   for(int dof=0; dof<num_trial_dof; ++dof) {
+                       dest(dof + vd * num_trial_dof) += y_e_2d(dof, vd);
+                   }
+               }
+            }
+         });
+
+      }, num_entities, thread_blocks, shmem_info.total_size, shmem_cache.ReadWrite());
+      
+      for (size_t i = 0; i < num_solutions; i++) {
+         const auto P_rt = get_restriction_transpose<entity_t>(sol_fields[i], element_dof_ordering, outputs /* doesn't matter for rt */); // wait, get_restriction_transpose needs a field operator! 
+         // wait, actually we can just apply get_restriction_transpose logic or use `output_restriction_transpose`.
+         // Wait! We can just use `restriction_transpose` directly.
+      }
+#else
+      MFEM_ABORT("MFEM must be built with Enzyme to use vjp.");
+#endif
+   });
+
    // Without this compile-time check, some valid instantiations of this method
    // will fail.
    if constexpr (derivative_ids_t::size() != 0)

fem/dfem/qfunction_apply.hpp

@@ -614,6 +614,251 @@ void apply_kernel_fwddiff_enzyme(
    process_qf_result(f_qp,
                      get<0>(fwddiff_apply_enzyme(qfunc, args, shadow_args, tuple<> {})));
 }
+template <typename func_t, typename ret_t, typename... arg_ts>
+MFEM_HOST_DEVICE inline
+void qfunction_wrapper_out(const func_t &f, ret_t &out, arg_ts &&...args)
+{
+   out = f(std::forward<arg_ts>(args)...);
+}
+
+template <typename T>
+MFEM_HOST_DEVICE inline
+void process_qf_result_rev(const DeviceTensor<1, real_t> &r, T &x)
+{
+   x = r(0);
+}
+
+template <typename T>
+MFEM_HOST_DEVICE inline
+void process_qf_result_rev(const DeviceTensor<1, real_t> &r, dual<T, T> &x)
+{
+   x.value = r(0);
+}
+
+template <typename T>
+MFEM_HOST_DEVICE inline
+void process_qf_result_rev(const DeviceTensor<1, real_t> &r, tensor<T> &x)
+{
+   x(0) = r(0);
+}
+
+template <typename T, int n>
+MFEM_HOST_DEVICE inline
+void process_qf_result_rev(const DeviceTensor<1, real_t> &r, tensor<T, n> &x)
+{
+   for (size_t i = 0; i < n; i++)
+   {
+      x(i) = r(i);
+   }
+}
+
+template <typename T, int n, int m>
+MFEM_HOST_DEVICE inline
+void process_qf_result_rev(const DeviceTensor<1, real_t> &r, tensor<T, n, m> &x)
+{
+   for (size_t i = 0; i < n; i++)
+   {
+      for (size_t j = 0; j < m; j++)
+      {
+         x(i, j) = r(i + n * j);
+      }
+   }
+}
+
+// In Reverse mode, enzyme expects interleaved dup arguments:
+// __enzyme_autodiff(..., enzyme_dup, &arg1, &shadow_arg1, enzyme_dup, &arg2, &shadow_arg2, ...)
+
+#define ENZYME_DUP_ARGS_1(A, SA) enzyme_dup, &(A), &(SA)
+#define ENZYME_DUP_ARGS_2(A, SA, B, SB) enzyme_dup, &(A), &(SA), enzyme_dup, &(B), &(SB)
+#define ENZYME_DUP_ARGS_3(A, SA, B, SB, C, SC) enzyme_dup, &(A), &(SA), enzyme_dup, &(B), &(SB), enzyme_dup, &(C), &(SC)
+#define ENZYME_DUP_ARGS_4(A, SA, B, SB, C, SC, D, SD) enzyme_dup, &(A), &(SA), enzyme_dup, &(B), &(SB), enzyme_dup, &(C), &(SC), enzyme_dup, &(D), &(SD)
+#define ENZYME_DUP_ARGS_5(A, SA, B, SB, C, SC, D, SD, E, SE) enzyme_dup, &(A), &(SA), enzyme_dup, &(B), &(SB), enzyme_dup, &(C), &(SC), enzyme_dup, &(D), &(SD), enzyme_dup, &(E), &(SE)
+
+template <typename qfunc_t, typename arg_ts, typename inactive_arg_ts>
+MFEM_HOST_DEVICE inline
+void autodiff_apply_enzyme_indexed_1(qfunc_t &qfunc, void *out, void *dout, arg_ts &args, arg_ts &shadow_args, inactive_arg_ts &inactive_args)
+{
+   using qf_return_t = typename create_function_signature<decltype(&qfunc_t::operator())>::type::return_t;
+   __enzyme_autodiff<void>((void*)&qfunction_wrapper_out<qfunc_t, qf_return_t, decltype(get<0>(args))>,
+      enzyme_const, &qfunc,
+      enzyme_dupnoneed, out, dout,
+      ENZYME_DUP_ARGS_1(get<0>(args), get<0>(shadow_args))
+   );
+}
+
+template <typename qfunc_t, typename arg_ts, typename inactive_arg_ts>
+MFEM_HOST_DEVICE inline
+void autodiff_apply_enzyme_indexed_2(qfunc_t &qfunc, void *out, void *dout, arg_ts &args, arg_ts &shadow_args, inactive_arg_ts &inactive_args)
+{
+   using qf_return_t = typename create_function_signature<decltype(&qfunc_t::operator())>::type::return_t;
+   __enzyme_autodiff<void>((void*)&qfunction_wrapper_out<qfunc_t, qf_return_t, decltype(get<0>(args)), decltype(get<1>(args))>,
+      enzyme_const, &qfunc,
+      enzyme_dupnoneed, out, dout,
+      ENZYME_DUP_ARGS_2(get<0>(args), get<0>(shadow_args), get<1>(args), get<1>(shadow_args))
+   );
+}
+
+template <typename qfunc_t, typename arg_ts, typename inactive_arg_ts>
+MFEM_HOST_DEVICE inline
+void autodiff_apply_enzyme_indexed_3(qfunc_t &qfunc, void *out, void *dout, arg_ts &args, arg_ts &shadow_args, inactive_arg_ts &inactive_args)
+{
+   using qf_return_t = typename create_function_signature<decltype(&qfunc_t::operator())>::type::return_t;
+   __enzyme_autodiff<void>((void*)&qfunction_wrapper_out<qfunc_t, qf_return_t, decltype(get<0>(args)), decltype(get<1>(args)), decltype(get<2>(args))>,
+      enzyme_const, &qfunc,
+      enzyme_dupnoneed, out, dout,
+      ENZYME_DUP_ARGS_3(get<0>(args), get<0>(shadow_args), get<1>(args), get<1>(shadow_args), get<2>(args), get<2>(shadow_args))
+   );
+}
+
+template <typename qfunc_t, typename arg_ts, typename inactive_arg_ts>
+MFEM_HOST_DEVICE inline
+void autodiff_apply_enzyme_indexed_4(qfunc_t &qfunc, void *out, void *dout, arg_ts &args, arg_ts &shadow_args, inactive_arg_ts &inactive_args)
+{
+   using qf_return_t = typename create_function_signature<decltype(&qfunc_t::operator())>::type::return_t;
+   __enzyme_autodiff<void>((void*)&qfunction_wrapper_out<qfunc_t, qf_return_t, decltype(get<0>(args)), decltype(get<1>(args)), decltype(get<2>(args)), decltype(get<3>(args))>,
+      enzyme_const, &qfunc,
+      enzyme_dupnoneed, out, dout,
+      ENZYME_DUP_ARGS_4(get<0>(args), get<0>(shadow_args), get<1>(args), get<1>(shadow_args), get<2>(args), get<2>(shadow_args), get<3>(args), get<3>(shadow_args))
+   );
+}
+
+template <typename qfunc_t, typename arg_ts, typename inactive_arg_ts>
+MFEM_HOST_DEVICE inline
+void autodiff_apply_enzyme_indexed_5(qfunc_t &qfunc, void *out, void *dout, arg_ts &args, arg_ts &shadow_args, inactive_arg_ts &inactive_args)
+{
+   using qf_return_t = typename create_function_signature<decltype(&qfunc_t::operator())>::type::return_t;
+   __enzyme_autodiff<void>((void*)&qfunction_wrapper_out<qfunc_t, qf_return_t, decltype(get<0>(args)), decltype(get<1>(args)), decltype(get<2>(args)), decltype(get<3>(args)), decltype(get<4>(args))>,
+      enzyme_const, &qfunc,
+      enzyme_dupnoneed, out, dout,
+      ENZYME_DUP_ARGS_5(get<0>(args), get<0>(shadow_args), get<1>(args), get<1>(shadow_args), get<2>(args), get<2>(shadow_args), get<3>(args), get<3>(shadow_args), get<4>(args), get<4>(shadow_args))
+   );
+}
+
+template <typename qfunc_t, typename arg_ts, typename inactive_arg_ts>
+MFEM_HOST_DEVICE inline
+void autodiff_apply_enzyme(qfunc_t &qfunc, void *out, void *dout, arg_ts &&args,
+                          arg_ts &&shadow_args,
+                          inactive_arg_ts &&inactive_args)
+{
+   constexpr size_t num_active_args = tuple_size<std::remove_reference_t<arg_ts>>::value;
+   
+   if constexpr (num_active_args == 1) {
+       autodiff_apply_enzyme_indexed_1(qfunc, out, dout, args, shadow_args, inactive_args);
+   } else if constexpr (num_active_args == 2) {
+       autodiff_apply_enzyme_indexed_2(qfunc, out, dout, args, shadow_args, inactive_args);
+   } else if constexpr (num_active_args == 3) {
+       autodiff_apply_enzyme_indexed_3(qfunc, out, dout, args, shadow_args, inactive_args);
+   } else if constexpr (num_active_args == 4) {
+       autodiff_apply_enzyme_indexed_4(qfunc, out, dout, args, shadow_args, inactive_args);
+   } else if constexpr (num_active_args == 5) {
+       autodiff_apply_enzyme_indexed_5(qfunc, out, dout, args, shadow_args, inactive_args);
+   } else {
+       MFEM_ABORT("Unsupported number of arguments for reverse mode AD.");
+   }
+}
+
+template <typename qfunc_t, typename arg_ts, size_t num_args>
+MFEM_HOST_DEVICE inline
+void apply_kernel_vjp_enzyme(
+   const DeviceTensor<1, real_t> &df_qp,
+   qfunc_t &qfunc,
+   arg_ts &args,
+   arg_ts &shadow_args,
+   const std::array<DeviceTensor<2>, num_args> &u,
+   const std::array<DeviceTensor<2>, num_args> &du,
+   int qp_idx)
+{
+   process_qf_args(u, args, qp_idx);
+   process_qf_args(du, shadow_args, qp_idx);
+   
+   using qf_return_t = typename create_function_signature<
+                       decltype(&qfunc_t::operator())>::type::return_t;
+   qf_return_t out{};
+   qf_return_t dout{};
+
+   process_qf_result_rev(df_qp, get<0>(dout));
+
+   autodiff_apply_enzyme(qfunc, &out, &dout, std::forward<arg_ts>(args), std::forward<arg_ts>(shadow_args), tuple<> {});
+}
+
+template <
+   typename qf_param_ts,
+   typename qfunc_t,
+   std::size_t num_fields>
+MFEM_HOST_DEVICE inline
+void call_qfunction_vjp(
+   qfunc_t &qfunc,
+   const std::array<DeviceTensor<2>, num_fields> &input_shmem,
+   const std::array<DeviceTensor<2>, num_fields> &shadow_shmem,
+   DeviceTensor<2> &residual_shmem_adj,
+   const int &rs_qp,
+   const int &num_qp,
+   const int &q1d,
+   const int &dimension,
+   const bool &use_sum_factorization)
+{
+   if (use_sum_factorization)
+   {
+      if (dimension == 1)
+      {
+         MFEM_FOREACH_THREAD_DIRECT(q, x, q1d)
+         {
+            auto qf_args = decay_tuple<qf_param_ts> {};
+            auto qf_shadow_args = decay_tuple<qf_param_ts> {};
+            auto r_adj = Reshape(&residual_shmem_adj(0, q), rs_qp);
+            apply_kernel_vjp_enzyme(r_adj, qfunc, qf_args, qf_shadow_args, input_shmem, shadow_shmem, q);
+         }
+      }
+      else if (dimension == 2)
+      {
+         MFEM_FOREACH_THREAD_DIRECT(qx, x, q1d)
+         {
+            MFEM_FOREACH_THREAD_DIRECT(qy, y, q1d)
+            {
+               const int q = qx + q1d * qy;
+               auto qf_args = decay_tuple<qf_param_ts> {};
+               auto qf_shadow_args = decay_tuple<qf_param_ts> {};
+               auto r_adj = Reshape(&residual_shmem_adj(0, q), rs_qp);
+               apply_kernel_vjp_enzyme(r_adj, qfunc, qf_args, qf_shadow_args, input_shmem, shadow_shmem, q);
+            }
+         }
+      }
+      else if (dimension == 3)
+      {
+         MFEM_FOREACH_THREAD_DIRECT(qx, x, q1d)
+         {
+            MFEM_FOREACH_THREAD_DIRECT(qy, y, q1d)
+            {
+               MFEM_FOREACH_THREAD_DIRECT(qz, z, q1d)
+               {
+                  const int q = qx + q1d * (qy + q1d * qz);
+                  auto qf_args = decay_tuple<qf_param_ts> {};
+                  auto qf_shadow_args = decay_tuple<qf_param_ts> {};
+                  auto r_adj = Reshape(&residual_shmem_adj(0, q), rs_qp);
+                  apply_kernel_vjp_enzyme(r_adj, qfunc, qf_args, qf_shadow_args, input_shmem, shadow_shmem, q);
+               }
+            }
+         }
+      }
+      else
+      {
+#if !(defined(MFEM_USE_CUDA) || defined(MFEM_USE_HIP))
+         MFEM_ABORT("unsupported dimension for sum factorization");
+#endif
+      }
+   }
+   else
+   {
+      MFEM_FOREACH_THREAD_DIRECT(q, x, num_qp)
+      {
+         auto qf_args = decay_tuple<qf_param_ts> {};
+         auto qf_shadow_args = decay_tuple<qf_param_ts> {};
+         auto r_adj = Reshape(&residual_shmem_adj(0, q), rs_qp);
+         apply_kernel_vjp_enzyme(r_adj, qfunc, qf_args, qf_shadow_args, input_shmem, shadow_shmem, q);
+      }
+   }
+   MFEM_SYNC_THREAD;
+}
 #endif // MFEM_USE_ENZYME
 
 } // namespace mfem::future

tests/unit/dfem/test_diffusion.cpp

@@ -227,6 +227,43 @@ void diffusion(const char *filename, int p)
       MPI_Barrier(MPI_COMM_WORLD);
    }
 
+   SECTION("vjp")
+   {
+#ifdef MFEM_USE_ENZYME
+      DOperator dop_mf(sol, {{Rho, &rho_ps}, {Coords, mfes}}, pmesh);
+      typename Diffusion<DIM>::MFApply mf_apply_qf;
+      auto derivatives = std::integer_sequence<size_t, U> {};
+      dop_mf.AddDomainIntegrator(mf_apply_qf,
+                                 tuple{ Gradient<U>{}, Identity<Rho>{},
+                                        Gradient<Coords>{}, Weight{} },
+                                 tuple{ Gradient<U>{} }, *ir,
+                                 all_domain_attr, derivatives);
+      dop_mf.SetParameters({ &rho_coeff_cv, nodes });
+      auto dRdU = dop_mf.GetDerivative(U, {&x}, {&rho_coeff_cv, nodes});
+
+      Vector V(pfes.GetTrueVSize()), VJP_expected(pfes.GetTrueVSize()), VJP_actual(pfes.GetTrueVSize());
+      V.Randomize(2);
+
+      pfes.GetRestrictionMatrix()->Mult(x, X);
+
+      // Compute vjp using DifferentiableOperator::vjp
+      dop_mf.vjp(X, V, VJP_actual);
+
+      // Compute vjp using DerivativeOperator::MultTranspose
+      dRdU->MultTranspose(V, VJP_expected);
+
+      VJP_actual -= VJP_expected;
+
+      real_t norm_global = 0.0;
+      real_t norm_local = VJP_actual.Normlinf();
+      MPI_Allreduce(&norm_local, &norm_global, 1, MPI_DOUBLE, MPI_MAX,
+                    pmesh.GetComm());
+
+      REQUIRE(norm_global == MFEM_Approx(0.0));
+      MPI_Barrier(MPI_COMM_WORLD);
+#endif
+   }
+
    SECTION("action vector")
    {
       ParFiniteElementSpace vpfes(&pmesh, &fec, DIM);