diff --git a/AD.jl b/AD.jl
new file mode 100644
index 000000000..107a23679
--- /dev/null
+++ b/AD.jl
@@ -0,0 +1,317 @@
+using GeoParams, ForwardDiff, StaticArrays, LinearAlgebra
+import GeoParams.MaterialParameters.ConstitutiveRelationships: compute_τII_parallel, compute_τII_harmonic, compute_εII_elements, compute_τII_elements
+
+# Define a range of rheological components
+v1 = SetDiffusionCreep("Dry Anorthite | Rybacki et al. (2006)")
+v2 = SetDislocationCreep("Dry Anorthite | Rybacki et al. (2006)")
+v3 = LinearViscous()
+v4 = LinearViscous(; η=1e22Pa * s)
+e1 = ConstantElasticity()           # elasticity
+e2 = SetConstantElasticity(; G=5e10, Kb=1e11)
+#pl1= DruckerPrager(C=1e6)                # plasticity
+pl1 = DruckerPrager(; C=1e6 / cosd(30))        # plasticity which ends up with the same yield stress as pl3
+pl2 = DruckerPrager(; C=1e6, ϕ=0, Ψ=10)      # plasticity
+pl3 = DruckerPrager(; C=1e6, ϕ=0)            # plasticity
+
+# Parallel elements
+p1 = Parallel(v3, v4)                # linear elements
+p2 = Parallel(v1, v2)                # includes nonlinear viscous elements
+p3 = Parallel(v1, v2, v3)             # includes nonlinear viscous elements
+p4 = Parallel(pl1, LinearViscous(; η=1e20Pa * s)) # viscoplastic regularisation
+
+# CompositeRheologies
+c1 = CompositeRheology(v1, v2)
+c2 = CompositeRheology(v3, v4)       # two linear rheologies
+c3 = CompositeRheology(v1, v2, e1)   # with elasticity
+c4 = CompositeRheology(v1, v3, p1)   # with linear || element
+c5 = CompositeRheology(v1, v4, p2)   # with nonlinear || element
+c6 = CompositeRheology(v1, v4, p1, p1) # with 2 || elements
+c7 = CompositeRheology(v4, e1)       # viscoelastic with linear viscosity
+c8 = CompositeRheology(v4, e1, pl1)   # with plastic element
+c9 = CompositeRheology(v4, e1, p4)    # with visco-plastic parallel element
+c10 = CompositeRheology(e1, pl3)      # elastoplastic
+c11 = CompositeRheology(e1, Parallel(pl3, LinearViscous(; η=1e19Pa * s)))      # elasto-viscoplastic
+c14 = CompositeRheology(SetConstantElasticity(G=1e10, Kb=1e10), LinearViscous(η=1e20), DruckerPrager(C=3e5, Ψ=1))   # case A
+
+
+# Check derivatives 
+vec1 = [c1 c2 c3 c4 c5 p1 p2 p3]
+args = (T=900.0, d=100e-6, τII_old=1e6, dt=1e8)
+εII, τII = 2e-15, 2e6
+
+v = c4   # problem with c3 (elasticity) and c4 (that has || elements) 
+
+compute_εII(c6, τII, args)
+compute_τII(c6, εII, args)
+
+ProfileCanvas.@profview for i in 1:1000000
+    compute_εII(v, τII, args)
+end
+
+@inline function derivative_all(f::F, x::R) where {F,R<:Real}
+    T = typeof(ForwardDiff.Tag(f, R))
+    res = f(ForwardDiff.Dual{T}(x, one(x)))
+    f, ∂f∂x = if res != 0.0
+        res.value, res.partials.values[1]
+    else
+        0.0, 0.0
+    end
+    return f, ∂f∂x
+end
+
+# im lazy so we define both functions here
+for tensor in (:εII, :τII)
+    fun1 = Symbol("λ_$(tensor)")
+    fun2 = Symbol("compute_$(tensor)")
+    @eval begin
+        function $(fun1)(vi, II, args)
+            # trick with nested lambdas so that it works with @generated
+            # It'd work with one lambda, but for some reason this is
+            # slightly more optimal
+            @inline λ(vi, II) = begin
+                _λ = II -> $(fun2)(vi, II, args)
+                derivative_all(_λ, II)
+            end
+        
+            λ(vi, II)
+        end
+    end
+end
+
+function λ_εII_nonplastic_AD(vi, τII, args)
+    # trick with nested lambdas so that it works with @generated
+    # It'd work with one lambda, but for some reason this is
+    # slightly more optimal
+    @inline λ(vi, τII) = begin
+        _λ = τII -> _compute_εII_nonplastic(vi, τII, args)
+        derivative_all(_λ, τII)
+    end
+
+    λ(vi, τII)
+end
+
+@generated function compute_II_AD_all(v::NTuple{N, AbstractConstitutiveLaw},
+    λ::F,
+    II,
+    args
+) where {N,F}
+    quote
+        Base.@_inline_meta
+        εII, dεII_dτII = 0.0, 0.0
+        Base.@nexprs $N i -> (
+            V = λ(v[i], II, args);
+            εII += V[1];
+            dεII_dτII += V[2];
+        )
+        return εII, dεII_dτII 
+    end
+end
+
+compute_εII_AD_all(v, τII, args) = compute_II_AD_all(v, λ_εII, τII, args)
+compute_εII_AD_all(v::Union{Parallel, CompositeRheology}, τII, args) = compute_εII_AD_all(v.elements, τII, args)
+
+compute_εII_AD_nonplastic_all(v, τII, args) = compute_II_AD_all(v, λ_εII_nonplastic_AD, τII, args)
+compute_εII_AD_nonplastic_all(v::Union{Parallel, CompositeRheology}, τII, args) = compute_εII_AD_nonplastic_all(v.elements, τII, args)
+
+compute_τII_AD_all(v, εII, args) = compute_II_AD_all(v, λ_τII, εII, args)
+compute_τII_AD_all(v::Union{Parallel, CompositeRheology}, εII, args) = compute_τII_AD_all(v.elements, εII, args)
+
+function compute_τII_AD(
+    v::CompositeRheology{T,N,
+        0,is_parallel,
+        0,is_plastic,
+        Nvol,is_vol,
+        false},
+    εII, 
+    args; 
+    tol=1e-6
+) where {N, T, is_parallel, is_plastic, Nvol, is_vol}
+
+    # Initial guess
+    τII = compute_τII_harmonic(v, εII, args)
+
+    # Local Iterations
+    iter = 0
+    ϵ = 2.0 * tol
+    τII_prev = τII
+    while ϵ > tol
+        iter += 1
+        # Newton scheme -> τII = τII - f(τII)/dfdτII. Therefore,
+        #   f(τII) = εII - strain_rate_circuit(v, τII, args) = 0
+        #   dfdτII = - dεII_dτII(v, τII, args) 
+        #   τII -= f / dfdτII 
+        εII_n, dεII_dτII_n = compute_εII_AD_all(v, τII, args)
+        τII = muladd(εII - εII_n, inv(dεII_dτII_n), τII)
+
+        ϵ = abs(τII - τII_prev) * inv(τII)
+        τII_prev = τII
+    end
+
+    return τII
+end
+
+@inline function local_iterations_τII_AD_all(
+    v::Parallel, τII, args; tol=1e-6
+)
+    # Initial guess
+    εII = compute_εII_harmonic(v, τII, args)
+
+    # Local Iterations
+    iter = 0
+    ϵ = 2.0 * tol
+    εII_prev = εII
+    while ϵ > tol
+        iter += 1
+        # Newton scheme -> τII = τII - f(τII)/dfdτII. Therefore,
+        #   f(τII) = εII - strain_rate_circuit(v, τII, args) = 0
+        #   dfdτII = - dεII_dτII(v, τII, args) 
+        #   τII -= f / dfdτII 
+        τII_n, dτII_dεII_n = compute_τII_AD_all(v, εII, args)
+        εII = muladd(τII - τII_n, inv(dτII_dεII_n), εII)
+
+        ϵ = abs(εII - εII_prev) * inv(εII)
+        εII_prev = εII
+    end
+    
+    return εII
+end
+
+## PROTOTYPES
+
+# # create a Static Vector of a vector of evaluated functions
+# @generated function SInput(funs::NTuple{N1, Function}, inputs::SVector{N2, T}) where {N1, N2, T}
+#     quote
+#         Base.Cartesian.@nexprs $N1 i -> f_i = funs[i](inputs...)
+#         Base.Cartesian.@ncall $N1 SVector{$N1, $T} f
+#     end
+# end
+
+
+# # create a Mutable Vector of a vector of evaluated functions
+# @generated function MInput(funs::NTuple{N1, Function}, inputs::MVector{N2, T}) where {N1, N2, T}
+#     quote
+#         Base.Cartesian.@nexprs $N1 i -> f_i = funs[i](inputs...)
+#         Base.Cartesian.@ncall  $N1 MVector{$N1, $T} f
+#     end
+# end
+
+SInput(funs::NTuple{N1, Function}, x::SVector{N2, T}) where {N1, N2, T} = SVector{N1,T}(funs[i](x) for i in 1:N1)
+# SInput(funs::NTuple{N1, Function}, inputs::SVector{N2, T}, args) where {N1, N2, T} = SVector{N1,T}(funs[i](inputs, args) for i in 1:N1)
+
+function SInput(funs::NTuple{N1, Function}, x::SVector{N2, T}, args) where {N1, N2, T} 
+    SVector{N1,T}(
+        funs[i](
+            i == 1 ?  SVector{2, T}(sum(x[i] for i in 1:N2-1), x[end]) : SVector{2, T}(x[1], x[end]), 
+            args
+        ) 
+        for i in 1:N1
+    )
+end
+
+MInput(funs::NTuple{N1, Function}, x::MVector{N2, T}) where {N1, N2, T} = MVector{N1,T}(funs[i](x) for i in 1:N1)
+MInput(funs::NTuple{N1, Function}, x::MVector{N2, T}, args) where {N1, N2, T} = MVector{N1,T}(funs[i](x, args) for i in 1:N1)
+
+
+@inline function jacobian(f, x::StaticArray)
+    T = typeof(ForwardDiff.Tag(f, eltype(x)))
+    result = ForwardDiff.static_dual_eval(T, f, x)
+    J = ForwardDiff.extract_jacobian(T, result, x)
+    f = extract_value(result)
+    return f, J
+end
+
+extract_value(result::SVector{N, ForwardDiff.Dual{Tag, T, N}}) where {N,T,Tag} = SVector{N,T}(result[i].value for i in 1:N)
+
+########################################################################################################################################## 
+##########################################################################################################################################
+c=c4
+
+# NOTE; not working for general composites, only for one single parallel element
+function compute_τII_par(
+    c::CompositeRheology{
+        T,    N,
+        Npar, is_par,
+        0,    is_plastic,
+        0,    is_vol
+    },
+    εII, 
+    args; 
+    tol = 1e-6
+) where {T, N, Npar, is_par, is_plastic, is_vol}
+    
+    # number of unknowns  
+    n = 1 + Npar; 
+
+    # initial guesses for stress and strain
+    εII_total = εII
+    εp_n = εII_total 
+    τ_n = compute_τII_harmonic(c, εII_total, args)
+
+    # x[1:end-1] = strain rate of parallel element, x[end] = total stress invariant
+    x = SVector{n,Float64}(i == 1 ? εp_n : τ_n for i in 1:n)
+
+    # define closures that will be differentiated
+    f1(x, args) = εII_total - compute_εII_elements(c, x[2], args) - x[1]
+    f2(x, args) = x[2] - compute_τII_parallel(c, x[1], args)
+    funs = f1, f2
+
+    # funs = ntuple(Val(n)) do i
+    #     f = if i ==1
+    #         (x, args) -> εII_total - compute_εII_elements(c, x[end], args) - sum(x[i] for i in 1:Npar)
+    #     else
+    #         (x, args) -> x[end] - compute_τII_parallel(c.elements, x[i-1], args)
+    #     end
+    # end
+
+    # Local Iterations
+    iter = 0
+    ϵ = 2 * tol
+    max_iter = 20
+    while (ϵ > tol) && (iter < max_iter)
+        iter += 1
+
+        x_n = x
+        args_n = merge(args, (τII = x[end], )) # we define a new arguments tuple to maintain type stability 
+        r, J = jacobian(input -> SInput(funs, input, args_n), x)
+        # update solution
+        dx  = J\r 
+        x  -= dx   
+        
+        ϵ = norm(@. abs(x-x_n))
+    end
+    return x[2], x[1]
+end
+compute_τII_par(c4, εII, args)
+
+@btime compute_τII_par($c, $εII, $args)
+
+compute_τII_par(c, εII, args)
+
+compute_εII(c, τII, args)
+compute_τII(c, εII, args)
+
+@btime compute_τII($v, $εII, $args)
+
+
+fp = ntuple()
+
+ntuple(Val(2)) do i
+    if i == 1
+        (x, args) -> εII_total - compute_εII_elements(c, x[2], args) - sum(x[i] for i in 1:(2-1))
+    else        
+        (x, args) -> x[i+1] - compute_τII_parallel(c, x[1], args)
+    end
+end
+
+macro foos(x, args)
+
+    esc(
+        quote
+            # Tuple(
+                (x,args) -> x[1] 
+            #     for i in length(x)
+            # )
+        end
+    )
+        
+end
diff --git a/src/AD_tools.jl b/src/AD_tools.jl
new file mode 100644
index 000000000..444662cf8
--- /dev/null
+++ b/src/AD_tools.jl
@@ -0,0 +1,19 @@
+# Hack into ForwardDiff.jl to return both the f'(x) and f(x)
+function derivative(f::F, x::R) where {F,R<:Real}
+    T = typeof(ForwardDiff.Tag(f, R))
+    res = f(ForwardDiff.Dual{T}(x, one(x)))
+    return res.value, res.partials.values[1]
+end
+
+# Return vector-evaluated function and its Jacobian  
+function jacobian(f, x::Abstractrray)
+    T = typeof(ForwardDiff.Tag(f, eltype(x)))
+    result = ForwardDiff.static_dual_eval(T, f, x)
+    J = ForwardDiff.extract_jacobian(T, result, x)
+    f = extract_value(result)
+    return f, J
+end
+
+extract_value(result::SVector{N, ForwardDiff.Dual{Tag, T, N}}) where {N,T,Tag} = SVector{N,T}(result[i].value for i in 1:N)
+extract_value(result::AbstractArray) = [result[i].value for i in eachindex]
+
diff --git a/src/CompositeRheologies/CompositeRheology.jl b/src/CompositeRheologies/CompositeRheology.jl
index 2eb5c061d..e3273e907 100644
--- a/src/CompositeRheologies/CompositeRheology.jl
+++ b/src/CompositeRheologies/CompositeRheology.jl
@@ -77,12 +77,12 @@ isvolumetricplastic(v::CompositeRheology{T, N,  Npar, is_parallel, Nplast, is_pl
 
 Computing `εII` as a function of `τII` for a composite element is the sum of the individual contributions
 """
-compute_εII(v::CompositeRheology{T,N}, τII::_T, args; tol=1e-6, verbose=false) where {T,_T,N} = nreduce(vi -> first(compute_εII(vi, τII, args)), v.elements)
-compute_εII(v::CompositeRheology{T,N}, τII::Quantity, args; tol=1e-6, verbose=false) where {T,N} = nreduce(vi -> first(compute_εII(vi, τII, args)), v.elements)
+compute_εII(v::CompositeRheology{T,N}, τII, args) where {T, N} = nreduce(vi -> first(compute_εII(vi, τII, args)), v.elements)
+compute_εII(v::CompositeRheology{T,N}, τII::Quantity, args) where {T,N} = nreduce(vi -> first(compute_εII(vi, τII, args)), v.elements)
 compute_εII(v::AbstractMaterialParamsStruct, τII, args) = compute_εII(v.CompositeRheology[1], τII, args)
 
 # As we don't do iterations, this is the same
-function compute_εII_AD(v::CompositeRheology, τII, args; tol=1e-6, verbose=false)
+function compute_εII_AD(v::CompositeRheology, τII, args)
     return  compute_εII(v, τII, args)
 end
 compute_εII_AD(v::AbstractMaterialParamsStruct, τII, args) = compute_εII_AD(v.CompositeRheology[1], τII, args)
@@ -93,8 +93,8 @@ compute_εII_AD(v::AbstractMaterialParamsStruct, τII, args) = compute_εII_AD(v
 
 Computing `εvol` as a function of `p` for a composite element is the sum of the individual contributions
 """
-compute_εvol(v::CompositeRheology{T,N}, p::_T, args; tol=1e-6, verbose=false) where {T,_T,N} = nreduce(vi -> first(compute_εvol(vi, p, args)), v.elements)
-compute_εvol(v::CompositeRheology{T,N}, p::Quantity, args; tol=1e-6, verbose=false) where {T,N} = nreduce(vi -> first(compute_εvol(vi, p, args)), v.elements)
+compute_εvol(v::CompositeRheology{T,N}, p, args) where {T, N} = nreduce(vi -> first(compute_εvol(vi, p, args)), v.elements)
+compute_εvol(v::CompositeRheology{T,N}, p::Quantity, args) where {T, N} = nreduce(vi -> first(compute_εvol(vi, p, args)), v.elements)
 
 
 # COMPUTE DEVIATORIC STRESS AND PRESSURE
@@ -274,23 +274,32 @@ dτII_dεII_AD(v::Union{Parallel,CompositeRheology}, εII, args) = ForwardDiff.d
 """
     compute_εII_nonplastic(v::CompositeRheology, TauII, args)
 
-Harmonic average of stress of all elements in a `CompositeRheology` structure that are not plastic elements
+Harmonic average of strain rate of all elements in a `CompositeRheology` structure that are not plastic elements
 """
 compute_εII_nonplastic(v::CompositeRheology{T,N}, τII::_T, args; tol=1e-6, verbose=false) where {T,_T,N} = nreduce(vi -> first(_compute_εII_nonplastic(vi, τII, args)), v.elements)
 _compute_εII_nonplastic(v, TauII, args) = first(compute_εII(v, TauII, args))
 _compute_εII_nonplastic(v::AbstractPlasticity, TauII, args) = 0.0
 
+"""
+    compute_εII_nonplastic(v::CompositeRheology, TauII, args)
+
+Harmonic average of strain rate of all elements in a `CompositeRheology` structure that are plastic elements
+"""
+compute_εII_plastic(v::CompositeRheology, τII, args, tol) = nreduce(vi -> _compute_εII_plastic(vi, τII, args, tol), v.elements)
+_compute_εII_plastic(v::AbstractPlasticity, τII, args, tol) = compute_εII(v, τII, args, tol)
+_compute_εII_plastic(v, τII, args, tol) = 0.0
+
 """
     compute_τII_harmonic(v::CompositeRheology, EpsII, args)
 
 Harmonic average of stress of all elements in a `CompositeRheology` structure that are not || elements
 """
 function compute_τII_harmonic(v::CompositeRheology{T,N}, EpsII::_T, args; tol=1e-6, verbose=false) where {T,_T,N} 
-    return inv(nreduce(vi -> first(_compute_τII_harmonic_element(vi, EpsII, args)), v.elements))
+    return 1.0/(nreduce(vi -> first(_compute_τII_harmonic_element(vi, EpsII, args)), v.elements))
 end
-_compute_τII_harmonic_element(v, EpsII, args) = inv(first(compute_τII(v, EpsII, args)))
-_compute_τII_harmonic_element(v::AbstractPlasticity, EpsII, args) = 0.0
-_compute_τII_harmonic_element(v::Parallel{T, N,  Nplast, is_plastic}, EpsII, args) where {T, N,  Nplast, is_plastic}  = 0.0
+@inline _compute_τII_harmonic_element(v, EpsII, args) = 1.0/(first(compute_τII(v, EpsII, args)))
+@inline _compute_τII_harmonic_element(v::AbstractPlasticity, EpsII, args) = 0.0
+@inline _compute_τII_harmonic_element(v::Parallel{T, N,  Nplast, is_plastic}, EpsII, args) where {T, N,  Nplast, is_plastic}  = 0.0
 
 """
     compute_p_harmonic(v::CompositeRheology, EpsVol, args)
@@ -341,19 +350,42 @@ end
 
 Sums the strainrate of all non-parallel and non-plastic elements in a `CompositeRheology` structure. Mostly internally used for jacobian iterations.
 """
-compute_εII_elements(v::CompositeRheology{T,N}, τII::_T, args; tol=1e-6, verbose=false) where {T,_T,N} = nreduce(vi -> first(_compute_εII_nonparallel(vi, τII, args)), v.elements)
-_compute_εII_nonparallel(v, TauII::_T, args) where {_T} = compute_εII(v, TauII, args)
-_compute_εII_nonparallel(v::Parallel, TauII::_T, args) where {_T} = zero(_T)
-_compute_εII_nonparallel(v::AbstractPlasticity, TauII::_T, args) where {_T} = zero(_T)
-
+compute_εII_elements(v::CompositeRheology{T,N}, τII, args) where {T,N} = nreduce(vi -> first(_compute_εII_nonparallel(vi, τII, args)), v.elements)
+_compute_εII_nonparallel(v, τII, args) = compute_εII(v, τII, args)
+_compute_εII_nonparallel(v::Parallel, τII::_T, args) where {_T} = zero(_T)
+_compute_εII_nonparallel(v::AbstractPlasticity, τII::_T, args) where {_T} = zero(_T)
 
 """
     compute_εvol_elements(v::CompositeRheology, TauII, args)
 
 Sums the strainrate of all non-parallel and non-plastic elements in a `CompositeRheology` structure. Mostly internally used for jacobian iterations.
 """
-compute_εvol_elements(v::CompositeRheology{T,N}, p::_T, args; tol=1e-6, verbose=false) where {T,_T,N} = nreduce(vi -> first(_compute_εvol_elements(vi, p, args)), v.elements)
-compute_εvol(v::Any, P::_T; kwargs...) where _T = _T(0) # in case nothing more specific is defined
+compute_εvol_elements(v::CompositeRheology{T,N}, p, args) where {T, N} = nreduce(vi -> first(_compute_εvol_elements(vi, p, args)), v.elements)
+compute_εvol(v::Any, P::_T; kwargs...) where _T =  zero(_T) # in case nothing more specific is defined
 _compute_εvol_elements(v, P::_T, args) where {_T} = compute_εvol(v, P, args)
 _compute_εvol_elements(v::Parallel, P::_T, args) where {_T} = zero(_T)
 _compute_εvol_elements(v::AbstractPlasticity, P::_T, args) where {_T} = zero(_T)
+
+"""
+    compute_τII_elements(v::CompositeRheology, TauII, args)
+
+Sums the stress of all non-parallel and non-plastic elements in a `CompositeRheology` structure. Mostly internally used for jacobian iterations.
+"""
+compute_τII_elements(v::CompositeRheology{T,N}, εII, args) where {T,N} = nreduce(vi -> first(_compute_τII_nonparallel(vi, εII, args)), v.elements)
+_compute_τII_nonparallel(v, εII, args) = compute_τII(v, εII, args)
+_compute_τII_nonparallel(v::Parallel, εII::_T, args) where {_T} = zero(_T)
+_compute_τII_nonparallel(v::AbstractPlasticity, εII::_T, args) where {_T} = zero(_T)
+
+
+"""
+    compute_τII_parallel(v::CompositeRheology, TauII, args)
+
+Sums the stress of all non-parallel and non-plastic elements in a `CompositeRheology` structure. Mostly internally used for jacobian iterations.
+"""
+compute_τII_parallel(v::CompositeRheology{T,N}, εII, args) where {T,N} = nreduce(vi -> first(_compute_τII_parallel(vi, εII, args)), v.elements)
+compute_τII_parallel(v::Parallel{T,N}, εII, args) where {T,N} = nreduce(vi -> first(_compute_τII_parallel(vi, εII, args)), v.elements)
+compute_τII_parallel(v, εII::_T, args) where {_T} = zero(_T)
+_compute_τII_parallel(v, εII::_T, args) where {_T} = zero(_T)
+_compute_τII_parallel(v::Parallel, εII, args) = compute_τII(v, εII, args)
+_compute_τII_parallel(v::AbstractPlasticity, εII::_T, args) where {_T} = zero(_T)
+
diff --git a/src/CompositeRheologies/NonlinearIterations.jl b/src/CompositeRheologies/NonlinearIterations.jl
index a25acf9a8..4c1f07b6e 100644
--- a/src/CompositeRheologies/NonlinearIterations.jl
+++ b/src/CompositeRheologies/NonlinearIterations.jl
@@ -65,61 +65,41 @@ Performs local iterations versus stress for a given strain rate using AD
     # Initial guess
     τII = compute_τII_harmonic(v, εII, args)
     
-    # @print(verbose, "initial stress_II = $τII")
-
-    # extract plastic element if it exists
-    v_pl = v[1]
-    if Nplast>0
-        for i=1:N
-            if is_plastic[i]
-                v_pl =  v[i]
-            end
-        end
-    end
-
     # Local Iterations
     iter = 0
     ϵ = 2.0 * tol
     τII_prev = τII
-    ε_pl = 0.0;
+    ε_pl = 0.0
+
     while (ϵ > tol) && (iter<10)
         iter += 1
-        #= 
-            Newton scheme -> τII = τII - f(τII)/dfdτII. 
-            Therefore,
-                f(τII) = εII - compute_εII(v, τII, args) = 0
-                dfdτII = - dεII_dτII(v, τII, args) 
-                τII -= f / dfdτII
-        =#
+        #  Newton scheme -> τII = τII - f(τII)/dfdτII. 
+        #  Therefore,
+        #      f(τII) = εII - compute_εII(v, τII, args) = 0
+        #      dfdτII = - dεII_dτII(v, τII, args) 
+        #      τII -= f / dfdτII
         
-        ε_np = compute_εII_nonplastic(v, τII, args)
-        dεII_dτII = dεII_dτII_nonplastic_AD(v, τII, args)
-
+        ε_np, dεII_dτII = compute_εII_AD_nonplastic_all(v, τII, args)
         f = εII - ε_np      # non-plastic contributions to residual
         
-        if Nplast>0
-           
+        if Nplast>0      
             # in case of plasticity, iterate for ε_pl
-            args = merge(args, (ε_np=ε_np,f=f))
-            ε_pl += compute_εII(v_pl, τII, args, tol=tol, verbose=verbose)
+            args = merge(args, (ε_np=ε_np,f=f))            
+            ε_pl += compute_εII_plastic(v, τII, args, tol)
 
             # add contributions to dεII_dτII:
-            if ε_pl>0.0
+            if ε_pl > 0.0
                 # in fact dε_pl/dτII = d(λ*∂Q∂τII(v_pl, τII))/dτII = 0 for DP
                 dεII_dτII += 0
             end
         end
         f -= ε_pl
 
-
         τII = muladd(f, inv(dεII_dτII), τII)
 
         ϵ = abs(τII - τII_prev) * inv(τII)
         τII_prev = τII
-        # @print(verbose, " iter $(iter) $ϵ τII=$τII")
     end
-    # @print(verbose, "final τII = $τII")
-    # @print(verbose, "---")
 
     return τII
 end
@@ -129,7 +109,7 @@ end
 
 Performs local iterations to compute the plastic strainrate. Note that the non-plastic strainrate, ε_np, should be part of `args`
 """
-function compute_εII(v::AbstractPlasticity, τII::_T, args; tol=1e-6, verbose=false) where _T
+function compute_εII(v::AbstractPlasticity, τII, args, tol)
 
     η_np  = (τII - args.τII_old)/(2.0*args.ε_np)
            
@@ -146,11 +126,12 @@ function compute_εII(v::AbstractPlasticity, τII::_T, args; tol=1e-6, verbose=f
         
         iter += 1
         τII_pl = τII -  2*η_np*λ*∂Q∂τII(v, τII_pl, args)       # update stress
-        F      = compute_yieldfunction(v, merge(args, (τII=τII_pl,)))
+        args   = merge(args, (τII=τII_pl,))
+        F      = compute_yieldfunction(v, args)
         
         dFdλ = ∂F∂τII(v, τII, args)*(2*η_np*∂Q∂τII(v, τII, args))
       
-        λ -= -F / dFdλ
+        λ -=  -F / dFdλ
 
         ϵ = F
 
@@ -162,7 +143,6 @@ function compute_εII(v::AbstractPlasticity, τII::_T, args; tol=1e-6, verbose=f
     return ε_pl
 end
 
-
 @inline function local_iterations_τII_AD(
     v::Parallel, τII::T, args; tol=1e-6, verbose=false
 ) where {T}
@@ -268,7 +248,7 @@ Performs local iterations versus strain rate for a given stress
     while ϵ > tol
         iter += 1
         f = τII - first(compute_τII(v, εII, args))
-        dfdεII = -dτII_dεII(v, εII, args)
+        dfdεII = - dτII_dεII(v, εII, args)
         εII -= f / dfdεII
 
         ϵ = abs(εII - εII_prev) / εII
@@ -287,7 +267,7 @@ end
 
 This performs nonlinear Newton iterations for `τII` with given `εII_total` for cases where we have both serial and parallel elements.
 """
-@inline function local_iterations_εII(
+function local_iterations_εII(
     c::CompositeRheology{T,
     N,
     Npar,is_par,
@@ -311,8 +291,7 @@ This performs nonlinear Newton iterations for `τII` with given `εII_total` for
 
     # @print(verbose,"τII guess = $τ_initial")
 
-    x    = @MVector ones(_T, n)
-    x   .= εII_total
+    x    = @MVector fill(εII_total, n)
     x[1] = τ_initial
 
     j = 1;
@@ -323,14 +302,12 @@ This performs nonlinear Newton iterations for `τII` with given `εII_total` for
         end
     end
     
-    r = @MVector zeros(_T,n);
-    J = @MMatrix zeros(_T, Npar+1,Npar+1)   # size depends on # of parallel objects (+ likely plastic elements)
+    r = @MVector zeros(_T,n)
+    J = @MMatrix zeros(_T, n, n)   # size depends on # of parallel objects (+ likely plastic elements)
     
     # Local Iterations
     iter = 0
     ϵ = 2 * tol
-    τII_prev = τ_initial
-    τ_parallel = _T(0)
     max_iter = 1000
     while (ϵ > tol) && (iter < max_iter)
         iter += 1
diff --git a/src/CompositeRheologies/Parallel.jl b/src/CompositeRheologies/Parallel.jl
index 749ed983b..e991621d2 100644
--- a/src/CompositeRheologies/Parallel.jl
+++ b/src/CompositeRheologies/Parallel.jl
@@ -4,20 +4,20 @@
 Put rheological elements in parallel 
 """
 struct Parallel{T, N,  Nplast, is_plastic, Nvol, is_vol} <: AbstractConstitutiveLaw{T}
-elements::T
+    elements::T
 end
 
 function Parallel(v::T) where T
-v           = tuple(v...)
-n           =   length(v)
+    v           = tuple(v...)
+    n           =   length(v)
 
-is_plastic = isa.(v,AbstractPlasticity)     # Is one of the elements a plastic element?
-Nplast = count(is_plastic)
+    is_plastic = isa.(v,AbstractPlasticity)     # Is one of the elements a plastic element?
+    Nplast = count(is_plastic)
 
-is_vol = isvolumetric.(v);
-Nvol   =   count(is_vol);
+    is_vol = isvolumetric.(v);
+    Nvol   =   count(is_vol);
 
-return Parallel{typeof(v),n, Nplast, is_plastic, Nvol, is_vol}(v)
+    return Parallel{typeof(v),n, Nplast, is_plastic, Nvol, is_vol}(v)
 end
 Parallel(a,b...) = Parallel((a,b...,)) 
 
diff --git a/src/CreepLaw/Elasticity_2.jl b/src/CreepLaw/Elasticity_2.jl
deleted file mode 100644
index c189feaa5..000000000
--- a/src/CreepLaw/Elasticity_2.jl
+++ /dev/null
@@ -1,74 +0,0 @@
-using ..MaterialParameters: MaterialParamsInfo
-import GeoParams.param_info
-
-export ConstantElasticity, dεII_dτII, dτII_dεII
-
-# ConstantElasticity  -------------------------------------------------------
-
-"""
-    ConstantElasticity(G, ν, K, E)
-
-Structure that holds parameters for constant, isotropic, linear elasticity.
-"""
-@with_kw_noshow struct ConstantElasticity{T,U,U1} <: AbstractCreepLaw{T}
-    G::GeoUnit{T,U} = 5e10Pa          # Elastic shear modulus
-    ν::GeoUnit{T,U1} = 0.5NoUnits      # Poisson ratio
-    K::GeoUnit{T,U} = 1e11Pa          # Elastic bulk modulus
-    E::GeoUnit{T,U} = 1e11Pa          # Elastic Young's modulus
-end
-ConstantElasticity(args...) = ConstantElasticity(convert.(GeoUnit, args)...)
-
-# Add multiple dispatch here to allow specifying combinations of 2 elastic parameters (say ν & E), to compute the others
-
-function param_info(s::ConstantElasticity) # info about the struct
-    return MaterialParamsInfo(; Equation=L"Constant elasticity")
-end
-
-# Calculation routines for linear viscous rheologies
-# All inputs must be non-dimensionalized (or converted to consitent units) GeoUnits
-@inline function computeCreepLaw_EpsII(
-    τII::_T, s::ConstantElasticity{_T}; τII_old::_T=zero(_T), dt::_T=one(_T), kwargs...
-) where {_T}
-    @unpack_val G = s
-    ε_el = (τII - τII_old) / (2.0 * G * dt)
-
-    return ε_el
-end
-
-@inline function dεII_dτII(τII, s::ConstantElasticity{_T}; dt::_T=1.0, kwargs...) where {_T}
-    @unpack_val G = s
-    return 1 / (2 * G * dt)
-end
-
-@inline function computeCreepLaw_TauII(
-    εII::_T, s::ConstantElasticity{_T}; dt::_T=one(_T), kwargs...
-) where {_T}
-    @unpack_val G = s
-    τII = G * dt
-
-    return τII
-end
-
-# """
-#     compute_elastic_shear_strainrate(s::ConstantElasticity{_T}; τII, τII_old, dt) 
-
-# Computes elastic strainrate given the deviatoric stress at the current (`τII`) and old timestep (`τII_old`), for a timestep `dt`:
-# ```math  
-#     \\dot{\\varepsilon}^{el} = {1 \\over 2 G} {D \\tau_{II} \\over Dt } ≈ {1 \\over 2 G} {\\tau_{II}- \\tilde{\\tau}_{II}^{old} \\over dt }
-# ```
-# Note that we here solve the scalar equation, which is sufficient for isotropic cases. In tensor form, it would be
-
-# ```math  
-#     {\\dot{\\varepsilon}^{el}}_{ij} = {1 \\over 2 G} { \\tau_{ij} - \\tilde{{\\tau_{ij}}}^{old} \\over dt }
-# ```
-# here ``\\tilde{{\\tau_{ij}}}^{old}`` is the rotated old deviatoric stress tensor to ensure objectivity (this can be done with Jaumann derivative, or also by using the full rotational formula).
-
-# """
-
-# Print info 
-function show(io::IO, g::ConstantElasticity)
-    return print(
-        io,
-        "Linear elasticity with shear modulus: G = $(UnitValue(g.G)), Poison's ratio: ν = $(UnitValue(g.ν)), bulk modulus: K = $(UnitValue(g.K)) and Young's module: E=$(UnitValue(g.E))",
-    )
-end
diff --git a/src/Elasticity/Elasticity.jl b/src/Elasticity/Elasticity.jl
index 5d55c98f2..301d199bc 100644
--- a/src/Elasticity/Elasticity.jl
+++ b/src/Elasticity/Elasticity.jl
@@ -196,8 +196,8 @@ Computes elastic volumetric strainrate given the pressure at the current (`P`) a
 
 """
 @inline function compute_εvol(
-    a::ConstantElasticity, P::_T; P_old=zero(precision(a)), dt=one(precision(a)), kwargs...
-) where {_T}
+    a::ConstantElasticity, P; P_old=zero(precision(a)), dt=one(precision(a)), kwargs...
+)
     @unpack_val Kb = a
     εvol_el = - (P - P_old) / (Kb * dt)
     return εvol_el
@@ -295,9 +295,9 @@ end
 
 # Single material phase
 @inline _elastic_ε(v::ConstantElasticity, τij_old, dt) = τij_old / (2 * v.G * dt)
-@inline _elastic_ε(v::Vararg{Any, N}) where {N} = 0.0
+@inline _elastic_ε(v::Vararg) = 0.0
 
-@inline effective_ε(εij::T, v, τij_old::T, dt) where {T} = εij + elastic_ε(v, τij_old, dt)
+@inline effective_ε(εij, v, τij_old, dt) = εij + elastic_ε(v, τij_old, dt)
 
 # Method for staggered grids
 @inline function effective_ε(
@@ -310,8 +310,8 @@ end
 end
 
 @inline function effective_ε(
-    εij::NTuple{N, T}, v, τij_old::NTuple{N,T}, dt
-) where {N,T}
+    εij::NTuple{N, T1}, v, τij_old::NTuple{N,T2}, dt
+) where {N,T1,T2}
     return ntuple(i -> effective_ε(εij[i], v, τij_old[i], dt), Val(N))
 end
 
@@ -389,7 +389,7 @@ end
 
 # Multiple material phases (collocated grid)
 
-@inline effective_ε(εij::T, v, τij_old::T, dt, phase::Int64) where {T} = εij + elastic_ε(v, τij_old, dt, phase)
+@inline effective_ε(εij, v, τij_old, dt, phase::Int64) = εij + elastic_ε(v, τij_old, dt, phase)
 
 # Method for staggered grids
 @inline function effective_ε(
@@ -402,8 +402,8 @@ end
 end
 
 @inline function effective_ε(
-    εij::NTuple{N, T}, v, τij_old::NTuple{N,T}, dt, phase::Int64
-) where {N,T}
+    εij::NTuple{N, Any}, v, τij_old::NTuple{N, Any}, dt, phase::Int64
+) where N
     return ntuple(i -> effective_ε(εij[i], v, τij_old[i], dt, phase), Val(N))
 end
 
diff --git a/src/Plasticity/DruckerPrager.jl b/src/Plasticity/DruckerPrager.jl
index 08450a56d..b14e95be5 100644
--- a/src/Plasticity/DruckerPrager.jl
+++ b/src/Plasticity/DruckerPrager.jl
@@ -50,9 +50,9 @@ function param_info(s::DruckerPrager) # info about the struct
 end
 
 # Calculation routines
-function (s::DruckerPrager{_T,U,U1})(;
-    P::_T=zero(_T), τII::_T=zero(_T), Pf::_T=zero(_T), kwargs...
-) where {_T,U,U1}
+@inline function (s::DruckerPrager)(;
+    P=zero(_T), τII=zero(_T), Pf=zero(_T), kwargs...
+)
     @unpack_val sinϕ, cosϕ, ϕ, C = s
     F = τII - cosϕ * C - sinϕ * (P - Pf)   # with fluid pressure (set to zero by default)
     return F
@@ -64,7 +64,7 @@ end
 Computes the plastic yield function `F` for a given second invariant of the deviatoric stress tensor `τII`,  `P` pressure, and `Pf` fluid pressure.
 """
 function compute_yieldfunction(
-    s::DruckerPrager{_T}; P::_T=zero(_T), τII::_T=zero(_T), Pf::_T=zero(_T)
+    s::DruckerPrager{_T}; P=zero(_T), τII=zero(_T), Pf=zero(_T)
 ) where {_T}
     return s(; P=P, τII=τII, Pf=Pf)
 end
@@ -94,12 +94,12 @@ end
 # Plastic Potential 
 
 # Derivatives w.r.t pressure
-∂Q∂P(p::DruckerPrager, P=zero(_T); τII=zero(_T), kwargs...) = -NumValue(p.sinΨ)
+∂Q∂P(p::DruckerPrager; kwargs...) = -NumValue(p.sinΨ)
 
 # Derivatives of yield function
-∂F∂τII(p::DruckerPrager, τII::_T; P=zero(_T), kwargs...) where _T  = _T(1)
-∂F∂P(p::DruckerPrager, P::_T; τII=zero(_T), kwargs...) where _T    = -NumValue(p.sinϕ)
-∂F∂λ(p::DruckerPrager, τII::_T; P=zero(_T), kwargs...) where _T    = _T(0) 
+∂F∂τII(p::DruckerPrager, τII::_T; kwargs...) where _T = one(_T)
+∂F∂P(p::DruckerPrager, P::_T; kwargs...) where _T = -NumValue(p.sinϕ)
+∂F∂λ(p::DruckerPrager, τII::_T;  kwargs...) where _T = zero(_T) 
 
 
 # Derivatives w.r.t stress tensor
@@ -108,34 +108,32 @@ end
 for t in (:NTuple,:SVector)
     @eval begin
         ## 3D derivatives 
-        ∂Q∂τxx(p::DruckerPrager, τij::$(t){6, T}) where T = 0.5 * τij[1] / second_invariant(τij)
-        ∂Q∂τyy(p::DruckerPrager, τij::$(t){6, T}) where T = 0.5 * τij[2] / second_invariant(τij)
-        ∂Q∂τzz(p::DruckerPrager, τij::$(t){6, T}) where T = 0.5 * τij[3] / second_invariant(τij)
-        ∂Q∂τyz(p::DruckerPrager, τij::$(t){6, T}) where T = τij[4] / second_invariant(τij)
-        ∂Q∂τxz(p::DruckerPrager, τij::$(t){6, T}) where T = τij[5] / second_invariant(τij)
-        ∂Q∂τxy(p::DruckerPrager, τij::$(t){6, T}) where T = τij[6] / second_invariant(τij) 
+        ∂Q∂τxx(::DruckerPrager, τij::$(t){6, T}) where T = 0.5 * τij[1] / second_invariant(τij)
+        ∂Q∂τyy(::DruckerPrager, τij::$(t){6, T}) where T = 0.5 * τij[2] / second_invariant(τij)
+        ∂Q∂τzz(::DruckerPrager, τij::$(t){6, T}) where T = 0.5 * τij[3] / second_invariant(τij)
+        ∂Q∂τyz(::DruckerPrager, τij::$(t){6, T}) where T = τij[4] / second_invariant(τij)
+        ∂Q∂τxz(::DruckerPrager, τij::$(t){6, T}) where T = τij[5] / second_invariant(τij)
+        ∂Q∂τxy(::DruckerPrager, τij::$(t){6, T}) where T = τij[6] / second_invariant(τij) 
         ## 2D derivatives 
-        ∂Q∂τxx(p::DruckerPrager, τij::$(t){3, T}) where T = 0.5 * τij[1] / second_invariant(τij)
-        ∂Q∂τyy(p::DruckerPrager, τij::$(t){3, T}) where T = 0.5 * τij[2] / second_invariant(τij)
-        ∂Q∂τxy(p::DruckerPrager, τij::$(t){3, T}) where T = τij[3] / second_invariant(τij) 
+        ∂Q∂τxx(::DruckerPrager, τij::$(t){3, T}) where T = 0.5 * τij[1] / second_invariant(τij)
+        ∂Q∂τyy(::DruckerPrager, τij::$(t){3, T}) where T = 0.5 * τij[2] / second_invariant(τij)
+        ∂Q∂τxy(::DruckerPrager, τij::$(t){3, T}) where T = τij[3] / second_invariant(τij) 
     end
 end
 
-∂Q∂τII(p::DruckerPrager, τII::_T; P=zero(_T), kwargs...) where _T = 0.5
+∂Q∂τII(p::DruckerPrager, τII; kwargs...) = 0.5
 
 """
     compute_εII(p::DruckerPrager{_T,U,U1}, λdot::_T, τII::_T,  P) 
 
 This computes plastic strain rate invariant for a given ``λdot``
 """
-function compute_εII(p::DruckerPrager{_T,U,U1}, λdot::_T, τII::_T, kwargs...) where {_T, U, U1}
+function compute_εII(p::DruckerPrager, λdot, τII, kwargs...)
     F = compute_yieldfunction(p, kwargs)
-    @show F, kwargs
-    if F>0
-        ε_pl = λdot*∂Q∂τII(p, τII)
-
+    ε_pl = if F > 0.0
+        λdot*∂Q∂τII(p, τII)
     else
-        ε_pl = 0.0
+        0.0
     end 
 
     return ε_pl
@@ -164,7 +162,7 @@ end
     `K` is the elastic bulk modulus, h is the harderning, and `τij`` is the stress tensor in Voigt notation.
     Equations from Duretz et al. 2019 G3
 """
-@inline function lambda(F::T, p::DruckerPrager, ηve::T, ηvp::T; K=zero(T), dt=zero(T), h=zero(T), τij=(one(T), one(T), one(T))) where T
+@inline function lambda(F, p::DruckerPrager{_T,U,U1}, ηve, ηvp; K=0.0, dt=0.0, h=0.0, τij=(1.0,1.0,1.0)) where {_T,U,U1}
     F * inv(ηve + ηvp + K * dt * p.sinΨ * p.sinϕ + h * p.cosϕ * plastic_strain(p, τij, zero(T)))
 end
 #-------------------------------------------------------------------------
\ No newline at end of file
diff --git a/src/Rheology_AD.jl b/src/Rheology_AD.jl
new file mode 100644
index 000000000..609187145
--- /dev/null
+++ b/src/Rheology_AD.jl
@@ -0,0 +1,384 @@
+# i'm lazy so let's define both functions programatically
+for tensor in (:εII, :τII)
+    fun1 = Symbol("λ_$(tensor)")
+    fun2 = Symbol("compute_$(tensor)")
+    @eval begin
+        # One can't define lambas inside @generated functions (i.e. as needed for FD.derivative)
+        # It'd work with one lambda, but for some obscure reason nested lambdas 
+        # yield a slightly more performant algorithm
+        function $(fun1)(vi, II, args)
+            # lambdas party
+            @inline λ(vi, II) = begin
+                _λ = II -> $(fun2)(vi, II, args)
+                derivative_all(_λ, II)
+            end
+        
+            λ(vi, II)
+        end
+    end
+end
+
+# evalute f(x) and f'(x) of second invariant of tensor A w.r.t. second invariant of tensor B
+@generated function compute_II_AD_all(v::NTuple{N, AbstractConstitutiveLaw}, λ::F, BII, args) where {N,F}
+    quote
+        Base.@_inline_meta
+        AII, dεII_dBII = 0.0, 0.0
+        Base.@nexprs $N i -> @inbounds ( 
+            V = λ(v[i], BII, args);
+            AII += V[1];
+            dεII_dBII += V[2];
+        )
+        return AII, dεII_dBII 
+    end
+end
+
+# thin wrappers to compute the second invariant of the strain rate using forward mode AD
+compute_εII_AD_all(v, τII, args) = compute_II_AD_all(v, λ_εII, τII, args)
+compute_εII_AD_all(v::Union{Parallel, CompositeRheology}, τII, args) = compute_εII_AD_all(v.elements, τII, args)
+
+# thin wrappers to compute the second invariant of the deviatoric stress using forward mode AD
+compute_τII_AD_all(v, εII, args) = compute_II_AD_all(v, λ_τII, εII, args)
+compute_τII_AD_all(v::Union{Parallel, CompositeRheology}, εII, args) = compute_τII_AD_all(v.elements, εII, args)
+
+# generic function to compute the second invariant of tensor B 
+# given the second invariant of tensor AII
+@inline function compute_II_AD(
+    v::CompositeRheology{T,N,
+        0,is_parallel,
+        0,is_plastic,
+        Nvol,is_vol,
+        false},
+    f_harmonic::F1,
+    f_tensor_AD::F2,
+    AII, 
+    args, 
+    tol
+) where {N, T, is_parallel, is_plastic, Nvol, is_vol, F1, F2}
+
+    # Initial guess
+    BII = f_harmonic(v, AII, args)
+
+    # Local Iterations
+    iter = 0
+    ϵ = 2.0 * tol
+    BII_prev = BII
+    while ϵ > tol
+        iter += 1
+        # Newton scheme -> τII = τII - f(τII)/dfdτII. Therefore,
+        #   f(τII) = εII - strain_rate_circuit(v, τII, args) = 0
+        #   dfdτII = - dεII_dτII(v, τII, args) 
+        #   τII -= f / dfdτII 
+        AII_n, dAII_dBII_n = f_tensor_AD(v, BII, args)
+        BII = muladd(AII - AII_n, inv(dAII_dBII_n), BII)
+
+        ϵ = abs(BII - BII_prev) * inv(BII)
+        BII_prev = BII
+    end
+
+    return BII
+end
+
+compute_τII_AD(v, εII, args; tol=1e-6) = compute_II_AD(v, compute_τII_harmonic, compute_εII_AD_all, εII, args, tol)
+compute_εII_AD(v, τII, args; tol=1e-6) = compute_II_AD(v, compute_εII_harmonic, compute_τII_AD_all, τII, args, tol)
+
+# generic function to compute the second invariant of tensor B 
+# given the second invariant of tensor AII
+@inline function local_iterations_II_AD(v::Parallel, f_harmonic::F1, f_tensor_AD::F2, AII, args, tol) where {F1, F2}
+    # Initial guess
+    BII = f_harmonic(v, AII, args)
+
+    # Local Iterations
+    iter = 0
+    ϵ = 2.0 * tol
+    BII_prev = BII
+    while ϵ > tol
+        iter += 1
+        # Newton scheme -> τII = τII - f(τII)/dfdτII. Therefore,
+        #   f(τII) = εII - strain_rate_circuit(v, τII, args) = 0
+        #   dfdτII = - dεII_dτII(v, τII, args) 
+        #   τII -= f / dfdτII 
+        AII_n, dAII_dBII_n = f_tensor_AD(v, BII, args)
+        BII = muladd(AII - AII_n, inv(dAII_dBII_n), BII)
+
+        ϵ = abs(BII - BII_prev) * inv(BII)
+        BII_prev = BII
+    end
+    
+    return BII
+end
+
+local_iterations_τII_AD(v, εII, args; tol=1e-6) = local_iterations_II_AD(v, compute_τII_harmonic, compute_εII_AD_all, εII, args, tol)
+local_iterations_εII_AD(v, τII, args; tol=1e-6) = local_iterations_II_AD(v, compute_εII_harmonic, compute_τII_AD_all, τII, args, tol)
+
+
+# PIRACY TIME FROM CompositeRheology.jl ------------------------------------------------------------
+
+# COMPUTE DEVIATORIC STRESS
+"""
+    τII = compute_τII_AD(v::CompositeRheology{T,N}, εII, args; tol=1e-6, verbose=false) 
+"""
+function compute_τII_AD(v::CompositeRheology, εII, args; tol=1e-6)
+    # A composite rheology case with parallel elements
+    τII = local_iterations_εII_AD(v, εII, args; tol=tol)
+    return τII
+end
+
+compute_τII(v::CompositeRheology{T,N,0}, εII, args; tol=1e-6) where {T,N} = compute_τII_AD(v, εII, args; tol=tol)
+
+"""
+    p,τII = compute_p_τII(v::CompositeRheology, εII, εvol, args; tol=1e-6, verbose=false) 
+
+This updates pressure `p` and deviatoric stress invariant `τII` in case the composite rheology structure has volumetric components, but does not contain plastic or parallel elements.
+The 'old' pressure should be stored in `args` as `args.P_old`   
+"""
+# TODO
+function compute_p_τII_AD(
+        v::CompositeRheology{T,N,
+                    Npar,is_parallel,
+                    Nplastic,is_plastic,
+                    Nvol,is_vol,
+                    false}, 
+        εII,
+        εvol,
+        args; 
+        tol=1e-6, verbose=false
+    ) where {T, N, Npar, is_parallel, Nplastic, is_plastic, Nvol, is_vol}
+
+    # A composite rheology case that may have volumetric elements, but the are not 
+    # tightly coupled, so we do NOT perform coupled iterations.
+    τII = local_iterations_εII_AD(v, εII, args; tol=tol)
+    P   = local_iterations_εvol_AD(v, εvol, args; tol=tol) # TODO
+
+    return P,τII
+end
+
+"""
+    p,τII = compute_p_τII(v::CompositeRheology, εII, εvol, args; tol=1e-6, verbose=false) 
+
+This updates pressure `p` and deviatoric stress invariant `τII` in case the composite rheology structure has no volumetric elemnts, but may contain plastic or parallel elements. 
+In that case, pressure is not updated (`args.P` is used instead).     
+"""
+function compute_p_τII_AD(
+        v::CompositeRheology{T,N,
+                    Npar,is_parallel,
+                    Nplast,is_plastic,
+                    0,is_vol, 
+                    false}, 
+        εII, 
+        εvol,
+        args;
+        tol=1e-6
+    ) where {T, N, _T, Npar, is_parallel, Nplast, is_plastic, is_vol}
+    
+    # A composite rheology case with no parallel element; iterations for τII
+    τII, = local_iterations_εII_AD(v, εII, args; tol=tol)
+    
+    P = any(keys(args) .=== :P_old) ? args.P_old : 0.0
+
+    return P, τII
+end
+
+"""
+    p,τII = compute_p_τII(v::CompositeRheology, εII, εvol, args; tol=1e-6, verbose=false) 
+
+This updates pressure `p` and deviatoric stress invariant `τII` in case the composite rheology structure has volumetric components and has volumetric plasticity
+The 'old' pressure should be stored in `args` as `args.P_old`   
+"""
+# TODO
+function compute_p_τII(
+        v::CompositeRheology{T,N,
+                        Npar,is_parallel,
+                        Nplastic,is_plastic,
+                        Nvol,is_vol,
+                        true}, 
+        εII::_T, 
+        εvol::_T,
+        args; 
+        tol=1e-6, verbose=false
+    ) where {T, N, _T, Npar, is_parallel, Nplastic, is_plastic, Nvol, is_vol}
+
+    # A composite rheology case that may have volumetric elements, but the are not 
+    # tightly coupled, so we do NOT perform coupled iterations.
+    out = local_iterations_εvol_εII_AD(v, εII, εvol, args; tol=tol, verbose=verbose) #TODO
+
+    τII = out[1]
+    P = out[end]
+
+    return P, τII, out[2:N]
+end
+
+
+"""
+    τII = local_iterations_εII_AD(v::CompositeRheology{T,N}, εII::_T, args; tol=1e-6, verbose=false)
+
+Performs local iterations versus stress for a given strain rate using AD
+"""
+@inline function local_iterations_εII_AD(
+    v::CompositeRheology{T,
+            N,
+            Npar,is_par,
+            Nplast,is_plastic,
+            Nvol,is_vol,
+            false},
+    εII::_T, 
+    args; 
+    tol=1e-6
+) where {N, T, _T, Npar, is_par, Nplast, is_plastic, Nvol, is_vol}
+    
+    # Initial guess
+    τII = compute_τII_harmonic(v, εII, args)
+    
+    # Local Iterations
+    iter = 0
+    ϵ = 2.0 * tol
+    τII_prev = τII
+    ε_pl = 0.0
+
+    while (ϵ > tol) && (iter<10)
+        iter += 1
+        #  Newton scheme -> τII = τII - f(τII)/dfdτII. 
+        #  Therefore,
+        #      f(τII) = εII - compute_εII(v, τII, args) = 0
+        #      dfdτII = - dεII_dτII(v, τII, args) 
+        #      τII -= f / dfdτII
+        
+        ε_np, dεII_dτII = compute_εII_AD_nonplastic_all(v, τII, args)
+        f = εII - ε_np      # non-plastic contributions to residual
+        
+        if Nplast>0      
+            # in case of plasticity, iterate for ε_pl
+            args = merge(args, (ε_np=ε_np,f=f))            
+            ε_pl += compute_εII_plastic(v, τII, args, tol)
+
+            # add contributions to dεII_dτII:
+            if ε_pl > 0.0
+                # in fact dε_pl/dτII = d(λ*∂Q∂τII(v_pl, τII))/dτII = 0 for DP
+                dεII_dτII += 0
+            end
+        end
+        f -= ε_pl
+
+        τII = muladd(f, inv(dεII_dτII), τII)
+
+        ϵ = abs(τII - τII_prev) * inv(τII)
+        τII_prev = τII
+    end
+
+    return τII
+end
+
+function local_iterations_τII_AD(v::Parallel, τII::T, args; tol=1e-6)
+    # Initial guess
+    εII = compute_εII_harmonic(v, τII, args)
+
+    # Local Iterations
+    iter = 0
+    ϵ = 2.0 * tol
+    εII_prev = εII
+    while ϵ > tol
+        iter += 1
+        εII_n, dτII_dεII_n = compute_εII_AD_all(v, εII, args)
+        εII = muladd(τII - εII_n, inv(dτII_dεII_n), εII)
+
+        ϵ = abs(εII - εII_prev) * inv(εII)
+        εII_prev = εII
+        
+    end
+    
+    return εII
+end
+
+# NOTE; not working for general composites, only for one single parallel element
+function compute_τII_par(
+    c::CompositeRheology{
+        T,    N,
+        Npar, is_par,
+        0,    is_plastic,
+        0,    is_vol
+    },
+    εII, 
+    args; 
+    tol = 1e-6
+) where {T, N, Npar, is_par, is_plastic, is_vol}
+    
+    # number of unknowns  
+    n = 1 + Npar; 
+
+    # initial guesses for stress and strain
+    εII_total = εII
+    εp_n = εII_total 
+    τ_n = compute_τII_harmonic(c, εII_total, args)
+
+    # x[1:end-1] = strain rate of parallel(s) element, x[end] = total stress invariant
+    x = SVector{n,Float64}(i != n ? εp_n : τ_n for i in 1:n)
+
+    # define closures to be differentiated
+    f1(x, args) = εII_total - compute_εII_elements(c, x[2], args) - x[1]
+    f2(x, args) = x[2] - compute_τII_parallel(c, x[1], args)
+    funs = f1, f2
+
+    # funs = ntuple(Val(n)) do i
+    #     f = if i ==1
+    #         (x, args) -> εII_total - compute_εII_elements(c, x[end], args) - sum(x[i] for i in 1:Npar)
+    #     else
+    #         (x, args) -> x[end] - compute_τII_parallel(c.elements, x[i-1], args)
+    #     end
+    # end
+
+    # Local Iterations
+    iter = 0
+    ϵ = 2 * tol
+    max_iter = 20
+    while (ϵ > tol) && (iter < max_iter)
+        iter += 1
+
+        x_n = x
+        args_n = merge(args, (τII = x[end], )) # we define a new arguments tuple to maintain type stability 
+        r, J = jacobian(input -> SInput(funs, input, args_n), x)
+
+        # update solution
+        dx  = J\r 
+        x  -= dx   
+        
+        ϵ = norm(@. abs(x-x_n))
+    end
+    
+    return x[end]
+
+end
+
+
+"""
+    p =local_iterations_εvol(v::CompositeRheology{T,N,0}, εvol::_T, args; tol=1e-6, verbose=false)
+
+Performs local iterations versus pressure for a given total volumetric strain rate for a given `CompositeRheology` element that does NOT include `Parallel` elements
+"""
+@inline function local_iterations_εvol(
+    v::CompositeRheology{T,N,
+                    0,is_parallel,
+                    0,is_plastic,
+                    Nvol,is_vol}, 
+    εvol,
+    args; 
+    tol=1e-6
+) where {N, T, is_parallel, is_plastic, Nvol, is_vol}
+
+    # Initial guess
+    p = compute_p_harmonic(v, εvol, args)
+    
+    # Local Iterations
+    iter = 0
+    ϵ = 2.0 * tol
+    p_prev = p
+    while ϵ > tol
+        iter += 1
+        εvol_n, dεvol_dp_n = derivative_all(p -> compute_εvol(v, p, args), p)
+        p = muladd(εvol - εvol_n, inv(dεvol_dp_n), p)
+
+        ϵ = abs(p - p_prev) * inv(abs(p))
+        p_prev = p
+
+    end
+    
+    return p
+end
\ No newline at end of file
diff --git a/src/StressComputations/StressComputations.jl b/src/StressComputations/StressComputations.jl
index 3b3183e94..f84f62746 100644
--- a/src/StressComputations/StressComputations.jl
+++ b/src/StressComputations/StressComputations.jl
@@ -111,11 +111,11 @@ This is for a collocated grid case with a single phase `phase`.
 """
 function compute_τij(
     v::NTuple{N1,AbstractMaterialParamsStruct},
-    εij::NTuple{N2,T},
+    εij::NTuple{N2,Any},
     args,
-    τij_old::NTuple{N2,T},
+    τij_old::NTuple{N2,Any},
     phase::I,
-) where {T,N1,N2,I<:Integer}
+) where {N1,N2,I<:Integer}
 
     # Second invariant of effective strainrate (taking elasticity into account)
     ε_eff = effective_ε(εij, v, τij_old, args.dt, phase)
@@ -505,15 +505,6 @@ end
 # ----------------------------------------------------------------------------------------
 
 ## Helper functions
-@inline function staggered_tensor_average(x::NTuple{N,Union{T,NTuple{4,T}}}) where {N,T}
-    ntuple(Val(N)) do i
-        Base.@_inline_meta
-        staggered_tensor_average(x[i])
-    end
-end
-
-staggered_tensor_average(x::NTuple{N,T}) where {N,T} = sum(x) / N
-staggered_tensor_average(x::T) where {T<:Number} = x
 
 @inline function volumetric_strainrate(x::NTuple{3,Union{T,NTuple{4,T}}}) where {T}
     return vol = x[1] + x[2]  #2D
@@ -522,3 +513,5 @@ end
 @inline function volumetric_strainrate(x::NTuple{6,Union{T,NTuple{4,T}}}) where {T}
     return vol = x[1] + x[2] + x[3] #3D
 end
+
+
diff --git a/src/TensorAlgebra/TensorAlgebra.jl b/src/TensorAlgebra/TensorAlgebra.jl
index 4d8ef799e..12b71e88b 100644
--- a/src/TensorAlgebra/TensorAlgebra.jl
+++ b/src/TensorAlgebra/TensorAlgebra.jl
@@ -5,9 +5,9 @@ import Base: (:)
 @inline average_pow2(x::NTuple{N,T}) where {N,T} = sum(xi^2 for xi in x) / N
 
 @inline function (:)(
-    A::Union{SVector{3,T},NTuple{3,T}}, B::Union{SVector{3,T},NTuple{3,T}}
-) where {T}
-    return (A[1] * B[1] + A[2] * B[2]) + T(2.0) * (A[3] * B[3])
+    A::Union{SVector{3,T1},NTuple{3,T1}}, B::Union{SVector{3,T2},NTuple{3,T2}}
+) where {T1, T2}
+    return (A[1] * B[1] + A[2] * B[2]) + 2.0 * (A[3] * B[3])
 end
 
 @inline function (:)(
@@ -20,11 +20,13 @@ end
 @inline (:)(A::SMatrix{M,M,T,N}, B::SMatrix{M,M,T,N}) where {M,N,T} = sum(A .* B)
 
 @inline second_invariant(A::NTuple{N,T}) where {N,T} = √(0.5 * (A:A))
-@inline second_invariant(A::Vararg{N,T}) where {N,T} = second_invariant(A)
+# @inline second_invariant(A::Vararg{N,T}) where {N,T} = second_invariant(A)
 @inline second_invariant(A::SMatrix) = √(0.5 * (A:A))
 @inline second_invariant(A::SVector) = √(0.5 * (A:A))
 @inline second_invariant(A::Matrix{T}) where {T} = √(0.5 * sum(Ai * Ai for Ai in A))
 
+@inline second_invariant(xx, yy, xy) = √( 0.5*(xx^2 + yy^2) + xy^2)
+
 """
     second_invariant_staggered(Aii::NTuple{2,T}, Axy::NTuple{4,T}) where {T} 
 
diff --git a/test/test_CompositeRheologies.jl b/test/test_CompositeRheologies.jl
index cb51d916a..1cd2406ca 100644
--- a/test/test_CompositeRheologies.jl
+++ b/test/test_CompositeRheologies.jl
@@ -1,6 +1,13 @@
 using Test
 using GeoParams, ForwardDiff
 
+using SnoopCompileCore
+invs = @snoopr using GeoParams
+tinf = @snoopi_deep SetDiffusionCreep("Dry Anorthite | Rybacki et al. (2006)");
+using SnoopCompile
+trees = invalidation_trees(invs);
+taletrees = precompile_blockers(trees, tinf)
+
 @testset "CompositeRheologies" begin
 
     # Define a range of rheological components