From 788090dcf05977024a1c408b4436c905e6d553d0 Mon Sep 17 00:00:00 2001 From: zjwegert <60646897+zjwegert@users.noreply.github.com> Date: Wed, 29 Apr 2026 11:30:25 +1000 Subject: [PATCH 1/5] Tag support Co-authored-by: Copilot --- src/Autodiff.jl | 24 ++++++++++++------------ test/AutodiffTests.jl | 16 ++++++++++++++++ 2 files changed, 28 insertions(+), 12 deletions(-) diff --git a/src/Autodiff.jl b/src/Autodiff.jl index 8d3a55db..c8367a87 100644 --- a/src/Autodiff.jl +++ b/src/Autodiff.jl @@ -112,9 +112,9 @@ end # autodiff_array_xxx function distributed_autodiff_array_gradient(a,i_to_x) - dummy_tag = ()->() + tag = x->ForwardDiff.gradient(a, x) 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 +127,9 @@ 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) 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 +147,9 @@ 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) 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 +163,9 @@ 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) 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 +213,9 @@ 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) 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 +245,9 @@ 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) 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) diff --git a/test/AutodiffTests.jl b/test/AutodiffTests.jl index f86a79d8..a13ac38c 100644 --- a/test/AutodiffTests.jl +++ b/test/AutodiffTests.jl @@ -43,6 +43,18 @@ 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))) + # Skeleton AD # I would like to compare the results, but we cannot be using FD in parallel... Λ = SkeletonTriangulation(model) @@ -213,4 +225,8 @@ function main(distribute,parts) skeleton_mf_different_fespace_trians(distribute,parts) end +with_mpi() do distribute + main(distribute, (2,2)) +end + end \ No newline at end of file From afdee07dfbdd25922d6e9b49082dcaa518bda466 Mon Sep 17 00:00:00 2001 From: zjwegert <60646897+zjwegert@users.noreply.github.com> Date: Wed, 29 Apr 2026 11:31:24 +1000 Subject: [PATCH 2/5] Minor --- test/AutodiffTests.jl | 4 ---- 1 file changed, 4 deletions(-) diff --git a/test/AutodiffTests.jl b/test/AutodiffTests.jl index a13ac38c..f4064a6b 100644 --- a/test/AutodiffTests.jl +++ b/test/AutodiffTests.jl @@ -225,8 +225,4 @@ function main(distribute,parts) skeleton_mf_different_fespace_trians(distribute,parts) end -with_mpi() do distribute - main(distribute, (2,2)) -end - end \ No newline at end of file From 2c00202b468831e577497485768433866e80f2ec Mon Sep 17 00:00:00 2001 From: zjwegert <60646897+zjwegert@users.noreply.github.com> Date: Wed, 29 Apr 2026 11:32:28 +1000 Subject: [PATCH 3/5] Update NEWS.md --- NEWS.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/NEWS.md b/NEWS.md index bf10ca45..4a69a35b 100644 --- a/NEWS.md +++ b/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). 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). From d3e0e1d243dc25cf1d220b43719d794535bd8625 Mon Sep 17 00:00:00 2001 From: zjwegert <60646897+zjwegert@users.noreply.github.com> Date: Mon, 4 May 2026 21:12:44 +1000 Subject: [PATCH 4/5] Conform with Gridap#1297 Co-authored-by: Copilot --- src/Autodiff.jl | 30 ++++++++++++++++++++++++++++-- test/AutodiffTests.jl | 22 ++++++++++++++++++++++ 2 files changed, 50 insertions(+), 2 deletions(-) diff --git a/src/Autodiff.jl b/src/Autodiff.jl index c8367a87..773cc6e1 100644 --- a/src/Autodiff.jl +++ b/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) @@ -113,6 +115,10 @@ end function distributed_autodiff_array_gradient(a,i_to_x) 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,tag),i_to_x) end @@ -128,6 +134,10 @@ end function distributed_autodiff_array_jacobian(a,i_to_x) 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,tag),i_to_x) end @@ -148,6 +158,10 @@ end function distributed_autodiff_array_gradient(a,i_to_x,j_to_i) 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,tag),i_to_x) end @@ -164,6 +178,10 @@ end function distributed_autodiff_array_jacobian(a,i_to_x,j_to_i) 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,tag),i_to_x) end @@ -214,6 +232,10 @@ end function distributed_autodiff_array_gradient(a, i_to_x, j_to_i::AbstractArray{<:SkeletonPair}) 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,tag),i_to_x) end @@ -246,6 +268,10 @@ end function distributed_autodiff_array_jacobian(a, i_to_x, j_to_i::AbstractArray{<:SkeletonPair}) 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,tag),i_to_x) end diff --git a/test/AutodiffTests.jl b/test/AutodiffTests.jl index f4064a6b..95da0ace 100644 --- a/test/AutodiffTests.jl +++ b/test/AutodiffTests.jl @@ -55,6 +55,28 @@ function main_sf(distribute,parts) ∂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) From 8fc3bf74e203abe9e75bf09f05ccf5c7d83dd099 Mon Sep 17 00:00:00 2001 From: zjwegert <60646897+zjwegert@users.noreply.github.com> Date: Mon, 4 May 2026 21:15:25 +1000 Subject: [PATCH 5/5] Update NEWS.md --- NEWS.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/NEWS.md b/NEWS.md index 4a69a35b..f8a21d10 100644 --- a/NEWS.md +++ b/NEWS.md @@ -7,7 +7,7 @@ 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). Since PR[#204](https://github.com/gridap/GridapDistributed.jl/pull/204). +- 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