AD tag support

NEWS.md

@@ -7,6 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+- Added tag support for AD to align with Gridap[#1181](https://github.com/gridap/Gridap.jl/pull/1181) and Gridap[#1297](https://github.com/gridap/Gridap.jl/pull/1297). Since PR[#204](https://github.com/gridap/GridapDistributed.jl/pull/204).
+
 ## [0.4.15] - 2026-03-23
 
 ### Fixed
@@ -15,7 +17,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [0.4.14] - 2026-03-20
 
-### Changed 
+### Changed
 
 - Added support for Gridap 0.20.0. Since PR[#200](https://github.com/gridap/GridapDistributed.jl/pull/200).
 

src/Autodiff.jl

@@ -11,11 +11,12 @@ end
 function FESpaces._gradient(f,uh,fuh::DistributedDomainContribution)
   local_terms = map(r -> DomainContribution(), get_parts(fuh))
   local_domains = tuple_of_arrays(map(Tuple∘get_domains,local_views(fuh)))
+  tag = x->Fields.gradient(f, uh)
   for local_trians in local_domains
     g = FESpaces._change_argument(gradient,f,local_trians,uh)
     cell_u = map(FESpaces.get_cell_dof_values,local_views(uh))
     cell_id = map(FESpaces._compute_cell_ids,local_views(uh),local_trians)
-    cell_grad = distributed_autodiff_array_gradient(g,cell_u,cell_id)
+    cell_grad = distributed_autodiff_array_gradient(g,cell_u,cell_id,tag)
     map(add_contribution!,local_terms,local_trians,cell_grad)
   end
   DistributedDomainContribution(local_terms)
@@ -29,11 +30,12 @@ end
 function FESpaces._jacobian(f,uh,fuh::DistributedDomainContribution)
   local_terms = map(r -> DomainContribution(), get_parts(fuh))
   local_domains = tuple_of_arrays(map(Tuple∘get_domains,local_views(fuh)))
+  tag = x->Fields.jacobian(f, uh)
   for local_trians in local_domains
     g = FESpaces._change_argument(jacobian,f,local_trians,uh)
     cell_u = map(FESpaces.get_cell_dof_values,local_views(uh))
     cell_id = map(FESpaces._compute_cell_ids,local_views(uh),local_trians)
-    cell_grad = distributed_autodiff_array_jacobian(g,cell_u,cell_id)
+    cell_grad = distributed_autodiff_array_jacobian(g,cell_u,cell_id,tag)
     map(add_contribution!,local_terms,local_trians,cell_grad)
   end
   DistributedDomainContribution(local_terms)
@@ -112,9 +114,13 @@ end
 # autodiff_array_xxx
 
 function distributed_autodiff_array_gradient(a,i_to_x)
-  dummy_tag = ()->()
+  tag = x->ForwardDiff.gradient(a, x)
+  distributed_autodiff_array_gradient(a,i_to_x,tag)
+end
+
+function distributed_autodiff_array_gradient(a,i_to_x,tag::Function)
   i_to_cfg = map(i_to_x) do i_to_x
-    lazy_map(ConfigMap(ForwardDiff.gradient,dummy_tag),i_to_x)
+    lazy_map(ConfigMap(ForwardDiff.gradient,tag),i_to_x)
   end
   i_to_xdual = map(i_to_cfg,i_to_x) do i_to_cfg, i_to_x
     lazy_map(DualizeMap(),i_to_cfg,i_to_x)
@@ -127,9 +133,13 @@ function distributed_autodiff_array_gradient(a,i_to_x)
 end
 
 function distributed_autodiff_array_jacobian(a,i_to_x)
-  dummy_tag = ()->()
+  tag = x->ForwardDiff.jacobian(a, x)
+  distributed_autodiff_array_jacobian(a,i_to_x,tag)
+end
+
+function distributed_autodiff_array_jacobian(a,i_to_x,tag::Function)
   i_to_cfg = map(i_to_x) do i_to_x
-    lazy_map(ConfigMap(ForwardDiff.jacobian,dummy_tag),i_to_x)
+    lazy_map(ConfigMap(ForwardDiff.jacobian,tag),i_to_x)
   end
   i_to_xdual = map(i_to_cfg,i_to_x) do i_to_cfg, i_to_x
     lazy_map(DualizeMap(),i_to_cfg,i_to_x)
@@ -147,9 +157,13 @@ function distributed_autodiff_array_hessian(a,i_to_x)
 end
 
 function distributed_autodiff_array_gradient(a,i_to_x,j_to_i)
-  dummy_tag = ()->()
+  tag = x->ForwardDiff.gradient(a, x)
+  distributed_autodiff_array_gradient(a,i_to_x,j_to_i,tag)
+end
+
+function distributed_autodiff_array_gradient(a,i_to_x,j_to_i,tag::Function)
   i_to_cfg = map(i_to_x) do i_to_x
-    lazy_map(ConfigMap(ForwardDiff.gradient,dummy_tag),i_to_x)
+    lazy_map(ConfigMap(ForwardDiff.gradient,tag),i_to_x)
   end
   i_to_xdual = map(i_to_cfg,i_to_x) do i_to_cfg, i_to_x
     lazy_map(DualizeMap(),i_to_cfg,i_to_x)
@@ -163,9 +177,13 @@ function distributed_autodiff_array_gradient(a,i_to_x,j_to_i)
 end
 
 function distributed_autodiff_array_jacobian(a,i_to_x,j_to_i)
-  dummy_tag = ()->()
+  tag = x->ForwardDiff.jacobian(a, x)
+  distributed_autodiff_array_jacobian(a,i_to_x,j_to_i,tag)
+end
+
+function distributed_autodiff_array_jacobian(a,i_to_x,j_to_i,tag::Function)
   i_to_cfg = map(i_to_x) do i_to_x
-    lazy_map(ConfigMap(ForwardDiff.jacobian,dummy_tag),i_to_x)
+    lazy_map(ConfigMap(ForwardDiff.jacobian,tag),i_to_x)
   end
   i_to_xdual = map(i_to_cfg,i_to_x) do i_to_cfg, i_to_x
     lazy_map(DualizeMap(),i_to_cfg,i_to_x)
@@ -213,9 +231,13 @@ function FESpaces._change_argument(op,f,local_trians::AbstractArray{<:SkeletonTr
 end
 
 function distributed_autodiff_array_gradient(a, i_to_x, j_to_i::AbstractArray{<:SkeletonPair})
-  dummy_tag = ()->()
+  tag = x->ForwardDiff.gradient(a, x)
+  distributed_autodiff_array_gradient(a, i_to_x, j_to_i, tag)
+end
+
+function distributed_autodiff_array_gradient(a, i_to_x, j_to_i::AbstractArray{<:SkeletonPair}, tag::Function)
   i_to_cfg = map(i_to_x) do i_to_x
-    lazy_map(ConfigMap(ForwardDiff.gradient,dummy_tag),i_to_x)
+    lazy_map(ConfigMap(ForwardDiff.gradient,tag),i_to_x)
   end
   i_to_xdual = map(i_to_cfg,i_to_x) do i_to_cfg, i_to_x
     lazy_map(DualizeMap(),i_to_cfg,i_to_x)
@@ -245,9 +267,13 @@ function distributed_autodiff_array_gradient(a, i_to_x, j_to_i::AbstractArray{<:
 end
 
 function distributed_autodiff_array_jacobian(a, i_to_x, j_to_i::AbstractArray{<:SkeletonPair})
-  dummy_tag = ()->()
+  tag = x->ForwardDiff.jacobian(a, x)
+  distributed_autodiff_array_jacobian(a, i_to_x, j_to_i, tag)
+end
+
+function distributed_autodiff_array_jacobian(a, i_to_x, j_to_i::AbstractArray{<:SkeletonPair}, tag::Function)
   i_to_cfg = map(i_to_x) do i_to_x
-    lazy_map(ConfigMap(ForwardDiff.jacobian,dummy_tag),i_to_x)
+    lazy_map(ConfigMap(ForwardDiff.jacobian,tag),i_to_x)
   end
   i_to_xdual = map(i_to_cfg,i_to_x) do i_to_cfg, i_to_x
     lazy_map(DualizeMap(),i_to_cfg,i_to_x)

test/AutodiffTests.jl

@@ -43,6 +43,40 @@ function main_sf(distribute,parts)
   b_AD = assemble_vector(gradient(g,uh),U)
   @test b ≈ b_AD
 
+  # Tags
+  dv = get_fe_basis(V)
+  dp = get_trial_fe_basis(U)
+  ph = interpolate(x->rand(),U)
+  ener(uh,ph) = ∫( 0.5*∇(uh)⋅∇(uh)*ph )*dΩ
+  res(uh,ph) = ∫( ∇(uh)⋅∇(dv)*ph )*dΩ
+  jac(uh,ph) = ∫( ∇(uh)⋅∇(dv)*dp )*dΩ
+
+  ∂2L∂u∂p_auto = assemble_matrix(jacobian(ph->gradient(uh->ener(uh,ph),uh),ph),U,U)
+  ∂2L∂u∂p = assemble_matrix(jac(uh,ph),U,U)
+  @test reduce(&,map(≈,partition(∂2L∂u∂p_auto),partition(∂2L∂u∂p)))
+
+  Γ_reg = BoundaryTriangulation(model)
+  Λ_reg = SkeletonTriangulation(model)
+  dΓ_reg = Measure(Γ_reg, 2)
+  dΛ_reg = Measure(Λ_reg, 2)
+
+  reffe_l2 = ReferenceFE(lagrangian, Float64, 1)
+  V_reg = TestFESpace(model, reffe_l2, conformity=:L2)
+  U_reg = TrialFESpace(V_reg)
+  dv_reg = get_fe_basis(V_reg)
+  dp_reg = get_trial_fe_basis(U_reg)
+  uh_reg = FEFunction(U_reg, rand(num_free_dofs(U_reg)))
+  ph_reg = FEFunction(U_reg, rand(num_free_dofs(U_reg)))
+
+  ener_reg(u, p) = ∫(0.5*u*u*p)*dΩ + ∫(0.5*u*u*p)*dΓ_reg + ∫(0.5*mean(u)*mean(u)*mean(p))*dΛ_reg
+  nested_ad_contrib = jacobian(p -> gradient(u -> ener_reg(u, p), uh_reg), ph_reg)
+  mat_nested_ad = assemble_matrix(nested_ad_contrib, U_reg, V_reg)
+  analytic_contrib = ∫(uh_reg*dv_reg*dp_reg)*dΩ + ∫(uh_reg*dv_reg*dp_reg)*dΓ_reg +
+                      ∫(mean(uh_reg)*mean(dv_reg)*mean(dp_reg))*dΛ_reg
+  mat_analytic = assemble_matrix(analytic_contrib, U_reg, V_reg)
+
+  @test reduce(&,map(≈,partition(mat_nested_ad),partition(mat_analytic)))
+
   # Skeleton AD
   # I would like to compare the results, but we cannot be using FD in parallel...
   Λ = SkeletonTriangulation(model)