diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
index 6df16f38..85c75900 100644
--- a/.github/workflows/CI.yml
+++ b/.github/workflows/CI.yml
@@ -6,6 +6,8 @@ on:
       - 'main'
       - 'release-'
     tags: '*'
+    paths-ignore:
+      - 'docs/**'
   pull_request:
   workflow_dispatch:
 
@@ -15,47 +17,26 @@ concurrency:
   cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }}
 
 jobs:
-  test:
-    name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.group }}
-    runs-on: ${{ matrix.os }}
+  tests:
+    name: "Tests"
     strategy:
       fail-fast: false
       matrix:
         version:
           - 'lts'
-          - '1' # automatically expands to the latest stable 1.x release of Julia
-        os:
-          - ubuntu-latest
-          - macOS-latest
-          - windows-latest
+          - '1'
         group:
           - ctmrg
           - boundarymps
           - examples
-    steps:
-      - uses: actions/checkout@v4
-      - uses: julia-actions/setup-julia@v2
-        with:
-          version: ${{ matrix.version }}
-      - uses: actions/cache@v4
-        env:
-          cache-name: cache-artifacts
-        with:
-          path: ~/.julia/artifacts
-          key: ${{ runner.os }}-test-${{ env.cache-name }}-${{ hashFiles('**/Project.toml') }}
-          restore-keys: |
-            ${{ runner.os }}-test-${{ env.cache-name }}-
-            ${{ runner.os }}-test-
-            ${{ runner.os }}-
-      - uses: julia-actions/julia-buildpkg@latest
-      - uses: julia-actions/julia-runtest@latest
-        env:
-          JULIA_NUM_THREADS: 4
-          GROUP: ${{ matrix.group }}
-      - uses: julia-actions/julia-processcoverage@v1
-      - uses: codecov/codecov-action@v4
-        if: always()
-        env:
-          CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
-        with:
-          file: lcov.info
+          - utility
+        os:
+          - ubuntu-latest
+          - macOS-latest
+          - windows-latest
+    uses: "QuantumKitHub/.github/.github/workflows/tests.yml@main"
+    with:
+      group: "${{ matrix.group }}"
+      julia-version: "${{ matrix.version }}"
+      os: "${{ matrix.os }}"
+    secrets: inherit
\ No newline at end of file
diff --git a/.github/workflows/FormatCheck.yml b/.github/workflows/FormatCheck.yml
index 1f0d0515..f58700c1 100644
--- a/.github/workflows/FormatCheck.yml
+++ b/.github/workflows/FormatCheck.yml
@@ -1,47 +1,14 @@
-name: FormatCheck
+name: "Format Check"
 
 on:
   push:
     branches:
       - 'main'
       - 'master'
-      - 'release-'
     tags: '*'
   pull_request:
 
 jobs:
-  build:
-    runs-on: ${{ matrix.os }}
-    strategy:
-      matrix:
-        version:
-          - '1' # automatically expands to the latest stable 1.x release of Julia
-        os:
-          - ubuntu-latest
-        arch:
-          - x64
-    steps:
-      - uses: julia-actions/setup-julia@latest
-        with:
-          version: ${{ matrix.version }}
-          arch: ${{ matrix.arch }}
-
-      - uses: actions/checkout@v4
-      - name: Install JuliaFormatter and format
-        # This will use the latest version by default but you can set the version like so:
-        #
-        # julia  -e 'using Pkg; Pkg.add(PackageSpec(name="JuliaFormatter", version="0.13.0"))'
-        run: |
-          julia  -e 'using Pkg; Pkg.add(PackageSpec(name="JuliaFormatter"))'
-          julia  -e 'using JuliaFormatter; format(".", verbose=true)'
-      - name: Format check
-        run: |
-          julia -e '
-          out = Cmd(`git diff --name-only`) |> read |> String
-          if out == ""
-              exit(0)
-          else
-              @error "Some files have not been formatted !!!"
-              write(stdout, out)
-              exit(1)
-          end'
\ No newline at end of file
+  format-check:
+    name: "Format Check"
+    uses: "QuantumKitHub/.github/.github/workflows/formatcheck.yml@main"
\ No newline at end of file
diff --git a/Project.toml b/Project.toml
index 2ebfd483..3c143ce1 100644
--- a/Project.toml
+++ b/Project.toml
@@ -13,6 +13,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LoggingExtras = "e6f89c97-d47a-5376-807f-9c37f3926c36"
 MPSKit = "bb1c41ca-d63c-52ed-829e-0820dda26502"
 MPSKitModels = "ca635005-6f8c-4cd1-b51d-8491250ef2ab"
+OhMyThreads = "67456a42-1dca-4109-a031-0a68de7e3ad5"
 OptimKit = "77e91f04-9b3b-57a6-a776-40b61faaebe0"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl
index f4bb9c83..c8a6b032 100644
--- a/src/PEPSKit.jl
+++ b/src/PEPSKit.jl
@@ -11,8 +11,10 @@ using LoggingExtras
 using MPSKit: loginit!, logiter!, logfinish!, logcancel!
 using MPSKitModels
 using FiniteDifferences
+using OhMyThreads
 
 include("utility/util.jl")
+include("utility/diffable_threads.jl")
 include("utility/svd.jl")
 include("utility/rotations.jl")
 include("utility/diffset.jl")
@@ -51,22 +53,45 @@ include("utility/symmetrization.jl")
     module Defaults
         const ctmrg_maxiter = 100
         const ctmrg_miniter = 4
-        const ctmrg_tol = 1e-12
-        const fpgrad_maxiter = 100
+        const ctmrg_tol = 1e-8
+        const fpgrad_maxiter = 30
         const fpgrad_tol = 1e-6
+        const ctmrgscheme = :simultaneous
+        const reuse_env = true
+        const trscheme = FixedSpaceTruncation()
+        const fwd_alg = TensorKit.SVD()
+        const rrule_alg = Arnoldi(; tol=1e-2fpgrad_tol, krylovdim=48, verbosity=-1)
+        const svd_alg = SVDAdjoint(; fwd_alg, rrule_alg)
+        const optimizer = LBFGS(32; maxiter=100, gradtol=1e-4, verbosity=2)
+        const gradient_linsolver = KrylovKit.BiCGStab(;
+            maxiter=Defaults.fpgrad_maxiter, tol=Defaults.fpgrad_tol
+        )
+        const iterscheme = :fixed
+        const gradient_alg = LinSolver(; solver=gradient_linsolver, iterscheme)
+        const scheduler = Ref{Scheduler}(Threads.nthreads() == 1 ? SerialScheduler() : DynamicScheduler())
     end
 
 Module containing default values that represent typical algorithm parameters.
 
-- `ctmrg_maxiter = 100`: Maximal number of CTMRG iterations per run
-- `ctmrg_miniter = 4`: Minimal number of CTMRG carried out
-- `ctmrg_tol = 1e-12`: Tolerance checking singular value and norm convergence
-- `fpgrad_maxiter = 100`: Maximal number of iterations for computing the CTMRG fixed-point gradient
-- `fpgrad_tol = 1e-6`: Convergence tolerance for the fixed-point gradient iteration
-- `rrule_alg = Arnoldi(; tol=1e1ctmrg_tol, krylovdim=48, verbosity=-1)`: Default cotangent linear problem algorithm
+- `ctmrg_maxiter`: Maximal number of CTMRG iterations per run
+- `ctmrg_miniter`: Minimal number of CTMRG carried out
+- `ctmrg_tol`: Tolerance checking singular value and norm convergence
+- `fpgrad_maxiter`: Maximal number of iterations for computing the CTMRG fixed-point gradient
+- `fpgrad_tol`: Convergence tolerance for the fixed-point gradient iteration
+- `ctmrgscheme`: Scheme for growing, projecting and renormalizing CTMRG environments
+- `reuse_env`: If `true`, the current optimization step is initialized on the previous environment
+- `trscheme`: Truncation scheme for SVDs and other decompositions
+- `fwd_alg`: SVD algorithm that is used in the forward pass
+- `rrule_alg`: Reverse-rule for differentiating that SVD
+- `svd_alg`: Combination of `fwd_alg` and `rrule_alg`
+- `optimizer`: Optimization algorithm for PEPS ground-state optimization
+- `gradient_linsolver`: Default linear solver for the `LinSolver` gradient algorithm
+- `iterscheme`: Scheme for differentiating one CTMRG iteration
+- `gradient_alg`: Algorithm to compute the gradient fixed-point
+- `scheduler`: Multi-threading scheduler which can be accessed via `set_scheduler!`
 """
 module Defaults
-    using TensorKit, KrylovKit, OptimKit
+    using TensorKit, KrylovKit, OptimKit, OhMyThreads
     using PEPSKit: LinSolver, FixedSpaceTruncation, SVDAdjoint
     const ctmrg_maxiter = 100
     const ctmrg_miniter = 4
@@ -77,7 +102,6 @@ module Defaults
     const sparse = false
     const reuse_env = true
     const trscheme = FixedSpaceTruncation()
-    const iterscheme = :fixed
     const fwd_alg = TensorKit.SDD()
     const rrule_alg = Arnoldi(; tol=1e-2fpgrad_tol, krylovdim=48, verbosity=-1)
     const svd_alg = SVDAdjoint(; fwd_alg, rrule_alg)
@@ -85,9 +109,56 @@ module Defaults
     const gradient_linsolver = KrylovKit.BiCGStab(;
         maxiter=Defaults.fpgrad_maxiter, tol=Defaults.fpgrad_tol
     )
+    const iterscheme = :fixed
     const gradient_alg = LinSolver(; solver=gradient_linsolver, iterscheme)
+
+    # OhMyThreads scheduler defaults
+    const scheduler = Ref{Scheduler}()
+    """
+        set_scheduler!([scheduler]; kwargs...)
+
+    Set `OhMyThreads` multi-threading scheduler parameters.
+
+    The function either accepts a `scheduler` as an `OhMyThreads.Scheduler` or
+    as a symbol where the corresponding parameters are specificed as keyword arguments.
+    For instance, a static scheduler that uses four tasks with chunking enabled
+    can be set via
+    ```
+    set_scheduler!(StaticScheduler(; ntasks=4, chunking=true))
+    ```
+    or equivalently with 
+    ```
+    set_scheduler!(:static; ntasks=4, chunking=true)
+    ```
+    For a detailed description of all schedulers and their keyword arguments consult the
+    [`OhMyThreads` documentation](https://juliafolds2.github.io/OhMyThreads.jl/stable/refs/api/#Schedulers).
+
+    If no `scheduler` is passed and only kwargs are provided, the `DynamicScheduler`
+    constructor is used with the provided kwargs.
+
+    To reset the scheduler to its default value, one calls `set_scheduler!` without passing
+    arguments which then uses the default `DynamicScheduler()`. If the number of used threads is
+    just one it falls back to `SerialScheduler()`.
+    """
+    function set_scheduler!(sc=OhMyThreads.Implementation.NotGiven(); kwargs...)
+        if isempty(kwargs) && sc isa OhMyThreads.Implementation.NotGiven
+            scheduler[] = Threads.nthreads() == 1 ? SerialScheduler() : DynamicScheduler()
+        else
+            scheduler[] = OhMyThreads.Implementation._scheduler_from_userinput(
+                sc; kwargs...
+            )
+        end
+        return nothing
+    end
+    export set_scheduler!
+
+    function __init__()
+        return set_scheduler!()
+    end
 end
 
+using .Defaults: set_scheduler!
+export set_scheduler!
 export SVDAdjoint, IterSVD, NonTruncSVDAdjoint
 export FixedSpaceTruncation, ProjectorAlg, CTMRG, CTMRGEnv, correlation_length
 export LocalOperator
diff --git a/src/algorithms/ctmrg/ctmrg.jl b/src/algorithms/ctmrg/ctmrg.jl
index 00fc4597..04b8651b 100644
--- a/src/algorithms/ctmrg/ctmrg.jl
+++ b/src/algorithms/ctmrg/ctmrg.jl
@@ -8,13 +8,14 @@ have different spaces, this truncation style is different from `TruncationSpace`
 struct FixedSpaceTruncation <: TensorKit.TruncationScheme end
 
 """
-    struct ProjectorAlg{S}(; svd_alg=Defaults.svd_alg, trscheme=Defaults.trscheme, verbosity=0)
-
-Algorithm struct collecting all projector related parameters.
-
-The `svd_alg` sets the SVD algorithm for decomposing the CTM environment. The truncation scheme
-has to be a `TensorKit.TruncationScheme`, and some SVD algorithms might have further restrictions
-on what kind of truncation scheme can be used.
+    struct ProjectorAlg{S}(; svd_alg=TensorKit.SVD(), trscheme=TensorKit.notrunc(),
+                           fixedspace=false, verbosity=0)
+
+Algorithm struct collecting all projector related parameters. The truncation scheme has to be
+a `TensorKit.TruncationScheme`, and some SVD algorithms might have further restrictions on what
+kind of truncation scheme can be used. If `fixedspace` is true, the truncation scheme is set to
+`truncspace(V)` where `V` is the environment bond space, adjusted to the corresponding
+environment direction/unit cell entry.
 """
 @kwdef struct ProjectorAlg{S<:SVDAdjoint,T}
     svd_alg::S = Defaults.svd_alg
@@ -177,13 +178,11 @@ ctmrg_logcancel!(log, iter, η, N) = @warnv 1 logcancel!(log, iter, η, N)
 
 """
     ctmrg_expand(state, envs, alg::CTMRG{M})
-    ctmrg_expand(dirs, state, envs::CTMRGEnv)
 
 Expand the environment by absorbing a new PEPS tensor.
 There are two modes of expansion: `M = :sequential` and `M = :simultaneous`.
 The first mode expands the environment in one direction at a time, for convenience towards
 the left. The second mode expands the environment in all four directions simultaneously.
-Alternatively, one can provide directly the `dirs` in which the environment is grown.
 """
 function ctmrg_expand(state, envs::CTMRGEnv, ::SequentialCTMRG)
     return ctmrg_expand([4, 1], state, envs)
@@ -191,20 +190,9 @@ end
 function ctmrg_expand(state, envs::CTMRGEnv, ::SimultaneousCTMRG)
     return ctmrg_expand(1:4, state, envs)
 end
-# function ctmrg_expand(dirs, state, envs::CTMRGEnv)  # TODO: This doesn't AD due to length(::Nothing)...
-#     drc_combinations = collect(Iterators.product(dirs, axes(state)...))
-#     return map(idx -> TensorMap(EnlargedCorner(state, envs, idx), idx[1]), drc_combinations)
-# end
-function ctmrg_expand(dirs, state, envs::CTMRGEnv{C,T}) where {C,T}
-    Qtype = tensormaptype(spacetype(C), 3, 3, storagetype(C))
-    Q = Zygote.Buffer(Array{Qtype,3}(undef, length(dirs), size(state)...))
-    dirs_enum = [(i, dir) for (i, dir) in enumerate(dirs)]
-    drc_combinations = collect(Iterators.product(dirs_enum, axes(state)...))
-    @fwdthreads for (d, r, c) in drc_combinations
-        ec = EnlargedCorner(state, envs, (d[2], r, c))
-        Q[d[1], r, c] = TensorMap(ec, d[2])  # Explicitly construct EnlargedCorner for now
-    end
-    return copy(Q)
+function ctmrg_expand(dirs, state, envs::CTMRGEnv)
+    coordinates = eachcoordinate(state, dirs)
+    return dtmap(idx -> TensorMap(EnlargedCorner(state, envs, idx), idx[1]), coordinates)
 end
 
 # ======================================================================================== #
@@ -221,15 +209,10 @@ function ctmrg_projectors(
     enlarged_envs, envs::CTMRGEnv{C,E}, alg::SequentialCTMRG
 ) where {C,E}
     projector_alg = alg.projector_alg
-    # pre-allocation
-    Prtype = tensormaptype(spacetype(E), numin(E), numout(E), storagetype(E))
-    P_bottom = Zygote.Buffer(envs.edges, axes(envs.corners, 2), axes(envs.corners, 3))
-    P_top = Zygote.Buffer(envs.edges, Prtype, axes(envs.corners, 2), axes(envs.corners, 3))
     ϵ = zero(real(scalartype(envs)))
 
-    directions = collect(Iterators.product(axes(envs.corners, 2), axes(envs.corners, 3)))
-    # @fwdthreads for (r, c) in directions
-    for (r, c) in directions
+    coordinates = eachcoordinate(envs)
+    projectors = dtmap(coordinates) do (r, c)
         # SVD half-infinite environment
         r′ = _prev(r, size(envs.corners, 2))
         QQ = halfinfinite_environment(enlarged_envs[1, r, c], enlarged_envs[2, r′, c])
@@ -248,26 +231,23 @@ function ctmrg_projectors(
         end
 
         # Compute projectors
-        P_bottom[r, c], P_top[r, c] = build_projectors(
-            U, S, V, enlarged_envs[1, r, c], enlarged_envs[2, r′, c]
-        )
+        return build_projectors(U, S, V, enlarged_envs[1, r, c], enlarged_envs[2, r′, c])
     end
 
-    return (copy(P_bottom), copy(P_top)), (; err=ϵ)
+    return (map(first, projectors), map(last, projectors)), (; err=ϵ)
 end
 function ctmrg_projectors(
     enlarged_envs, envs::CTMRGEnv{C,E}, alg::SimultaneousCTMRG
 ) where {C,E}
     projector_alg = alg.projector_alg
     # pre-allocation
-    P_left, P_right = Zygote.Buffer.(projector_type(envs.edges))
     U, V = Zygote.Buffer.(projector_type(envs.edges))
     # Corner type but with real numbers
     S = Zygote.Buffer(U.data, tensormaptype(spacetype(C), 1, 1, real(scalartype(E))))
 
     ϵ = zero(real(scalartype(envs)))
-    drc_combinations = collect(Iterators.product(axes(envs.corners)...))
-    @fwdthreads for (dir, r, c) in drc_combinations
+    coordinates = eachcoordinate(envs, 1:4)
+    projectors = dtmap(coordinates) do (dir, r, c)
         # Row-column index of next enlarged corner
         next_rc = if dir == 1
             (r, _next(c, size(envs.corners, 3)))
@@ -301,7 +281,7 @@ function ctmrg_projectors(
         end
 
         # Compute projectors
-        P_left[dir, r, c], P_right[dir, r, c] = build_projectors(
+        return build_projectors(
             U_local,
             S_local,
             V_local,
@@ -310,7 +290,9 @@ function ctmrg_projectors(
         )
     end
 
-    return (copy(P_left), copy(P_right)), (; err=ϵ, U=copy(U), S=copy(S), V=copy(V))
+    P_left = map(first, projectors)
+    P_right = map(last, projectors)
+    return (P_left, P_right), (; err=ϵ, U=copy(U), S=copy(S), V=copy(V))
 end
 
 """
@@ -374,9 +356,7 @@ function ctmrg_renormalize(projectors, state, envs, ::SequentialCTMRG)
     end
 
     # Apply projectors to renormalize corners and edges
-    coordinates = collect(Iterators.product(axes(state)...))
-    # @fwdthreads for (r, c) in coordinates
-    for (r, c) in coordinates
+    for (r, c) in eachcoordinate(state)
         C_southwest = renormalize_bottom_corner((r, c), envs, projectors)
         corners[SOUTHWEST, r, c] = C_southwest / norm(C_southwest)
 
@@ -390,12 +370,9 @@ function ctmrg_renormalize(projectors, state, envs, ::SequentialCTMRG)
     return CTMRGEnv(copy(corners), copy(edges))
 end
 function ctmrg_renormalize(enlarged_envs, projectors, state, envs, ::SimultaneousCTMRG)
-    corners = Zygote.Buffer(envs.corners)
-    edges = Zygote.Buffer(envs.edges)
     P_left, P_right = projectors
-
-    drc_combinations = collect(Iterators.product(axes(envs.corners)...))
-    @fwdthreads for (dir, r, c) in drc_combinations
+    coordinates = eachcoordinate(envs, 1:4)
+    corners_edges = dtmap(coordinates) do (dir, r, c)
         if dir == NORTH
             corner = renormalize_northwest_corner((r, c), enlarged_envs, P_left, P_right)
             edge = renormalize_north_edge((r, c), envs, P_left, P_right, state)
@@ -409,9 +386,8 @@ function ctmrg_renormalize(enlarged_envs, projectors, state, envs, ::Simultaneou
             corner = renormalize_southwest_corner((r, c), enlarged_envs, P_left, P_right)
             edge = renormalize_west_edge((r, c), envs, P_left, P_right, state)
         end
-        corners[dir, r, c] = corner / norm(corner)
-        edges[dir, r, c] = edge / norm(edge)
+        return corner / norm(corner), edge / norm(edge)
     end
 
-    return CTMRGEnv(copy(corners), copy(edges))
+    return CTMRGEnv(map(first, corners_edges), map(last, corners_edges))
 end
diff --git a/src/algorithms/ctmrg/gaugefix.jl b/src/algorithms/ctmrg/gaugefix.jl
index 849a8a89..5e63d5ef 100644
--- a/src/algorithms/ctmrg/gaugefix.jl
+++ b/src/algorithms/ctmrg/gaugefix.jl
@@ -8,14 +8,14 @@ element-wise converged to `envprev`.
 """
 function gauge_fix(envprev::CTMRGEnv{C,T}, envfinal::CTMRGEnv{C,T}) where {C,T}
     # Check if spaces in envprev and envfinal are the same
-    same_spaces = map(Iterators.product(axes(envfinal.edges)...)) do (dir, r, c)
+    same_spaces = map(eachcoordinate(envfinal, 1:4)) do (dir, r, c)
         space(envfinal.edges[dir, r, c]) == space(envprev.edges[dir, r, c]) &&
             space(envfinal.corners[dir, r, c]) == space(envprev.corners[dir, r, c])
     end
     @assert all(same_spaces) "Spaces of envprev and envfinal are not the same"
 
     # Try the "general" algorithm from https://arxiv.org/abs/2311.11894
-    signs = map(Iterators.product(axes(envfinal.edges)...)) do (dir, r, c)
+    signs = map(eachcoordinate(envfinal, 1:4)) do (dir, r, c)
         # Gather edge tensors and pretend they're InfiniteMPSs
         if dir == NORTH
             Tsprev = circshift(envprev.edges[dir, r, :], 1 - c)
@@ -70,7 +70,7 @@ end
 
 # Explicit fixing of relative phases (doing this compactly in a loop is annoying)
 function fix_relative_phases(envfinal::CTMRGEnv, signs)
-    corners_fixed = map(Iterators.product(axes(envfinal.corners)...)) do (dir, r, c)
+    corners_fixed = map(eachcoordinate(envfinal, 1:4)) do (dir, r, c)
         if dir == NORTHWEST
             fix_gauge_northwest_corner((r, c), envfinal, signs)
         elseif dir == NORTHEAST
@@ -82,7 +82,7 @@ function fix_relative_phases(envfinal::CTMRGEnv, signs)
         end
     end
 
-    edges_fixed = map(Iterators.product(axes(envfinal.corners)...)) do (dir, r, c)
+    edges_fixed = map(eachcoordinate(envfinal, 1:4)) do (dir, r, c)
         if dir == NORTHWEST
             fix_gauge_north_edge((r, c), envfinal, signs)
         elseif dir == NORTHEAST
@@ -99,7 +99,8 @@ end
 function fix_relative_phases(
     U::Array{Ut,3}, V::Array{Vt,3}, signs
 ) where {Ut<:AbstractTensorMap,Vt<:AbstractTensorMap}
-    U_fixed = map(Iterators.product(axes(U)...)) do (dir, r, c)
+    U_fixed = map(CartesianIndices(U)) do I
+        dir, r, c = I.I
         if dir == NORTHWEST
             fix_gauge_north_left_vecs((r, c), U, signs)
         elseif dir == NORTHEAST
@@ -111,7 +112,8 @@ function fix_relative_phases(
         end
     end
 
-    V_fixed = map(Iterators.product(axes(V)...)) do (dir, r, c)
+    V_fixed = map(CartesianIndices(V)) do I
+        dir, r, c = I.I
         if dir == NORTHWEST
             fix_gauge_north_right_vecs((r, c), V, signs)
         elseif dir == NORTHEAST
@@ -191,7 +193,8 @@ function calc_elementwise_convergence(envfinal::CTMRGEnv, envfix::CTMRGEnv; atol
     @debug "maxᵢⱼ|Tⁿ⁺¹ - Tⁿ|ᵢⱼ = $ΔTmax   mean |Tⁿ⁺¹ - Tⁿ|ᵢⱼ = $ΔTmean"
 
     # Check differences for all tensors in unit cell to debug properly
-    for (dir, r, c) in Iterators.product(axes(envfinal.edges)...)
+    for I in CartesianIndices(ΔT)
+        dir, r, c = I.I
         @debug(
             "$((dir, r, c)): all |Cⁿ⁺¹ - Cⁿ|ᵢⱼ < ϵ: ",
             all(x -> abs(x) < atol, convert(Array, ΔC[dir, r, c])),
diff --git a/src/algorithms/toolbox.jl b/src/algorithms/toolbox.jl
index 4f48cdf4..13e725e3 100644
--- a/src/algorithms/toolbox.jl
+++ b/src/algorithms/toolbox.jl
@@ -1,9 +1,10 @@
 function MPSKit.expectation_value(peps::InfinitePEPS, O::LocalOperator, envs::CTMRGEnv)
     checklattice(peps, O)
-    return sum(O.terms) do (inds, operator)  # TODO: parallelize this map, especially for the backwards pass
+    term_vals = dtmap([O.terms...]) do (inds, operator)  # OhMyThreads can't iterate over O.terms directly
         contract_localoperator(inds, operator, peps, peps, envs) /
         contract_localnorm(inds, peps, peps, envs)
     end
+    return sum(term_vals)
 end
 
 function costfun(peps::InfinitePEPS, envs::CTMRGEnv, O::LocalOperator)
diff --git a/src/environments/ctmrg_environments.jl b/src/environments/ctmrg_environments.jl
index cf90e2f0..460faf5b 100644
--- a/src/environments/ctmrg_environments.jl
+++ b/src/environments/ctmrg_environments.jl
@@ -105,7 +105,8 @@ function CTMRGEnv(
     corners = Array{C_type}(undef, 4, size(Ds_north)...)
     edges = Array{T_type}(undef, 4, size(Ds_north)...)
 
-    for (r, c) in Iterators.product(axes(Ds_north)...)
+    for I in CartesianIndices(Ds_north)
+        r, c = I.I
         edges[NORTH, r, c] = _edge_tensor(
             f,
             T,
@@ -371,6 +372,13 @@ function Base.rotl90(env::CTMRGEnv{C,T}) where {C,T}
 end
 
 Base.eltype(env::CTMRGEnv) = eltype(env.corners[1])
+Base.axes(x::CTMRGEnv, args...) = axes(x.corners, args...)
+function eachcoordinate(x::CTMRGEnv)
+    return collect(Iterators.product(axes(x, 2), axes(x, 3)))
+end
+function eachcoordinate(x::CTMRGEnv, dirs)
+    return collect(Iterators.product(dirs, axes(x, 2), axes(x, 3)))
+end
 
 # In-place update of environment
 function update!(env::CTMRGEnv{C,T}, env´::CTMRGEnv{C,T}) where {C,T}
diff --git a/src/operators/infinitepepo.jl b/src/operators/infinitepepo.jl
index a33c7e2c..e136aa70 100644
--- a/src/operators/infinitepepo.jl
+++ b/src/operators/infinitepepo.jl
@@ -125,7 +125,8 @@ function initializePEPS(
     T::InfinitePEPO{<:PEPOTensor{S}}, vspace::S
 ) where {S<:ElementarySpace}
     Pspaces = Array{S,2}(undef, size(T, 1), size(T, 2))
-    for (i, j) in product(1:size(T, 1), 1:size(T, 2))
+    for i in axes(T, 1)
+        j in axes(T, 2)
         Pspaces[i, j] = space(T, i, j)
     end
     Nspaces = repeat([vspace], size(T, 1), size(T, 2))
diff --git a/src/states/infinitepeps.jl b/src/states/infinitepeps.jl
index a57af615..51b353d4 100644
--- a/src/states/infinitepeps.jl
+++ b/src/states/infinitepeps.jl
@@ -124,6 +124,12 @@ Base.repeat(T::InfinitePEPS, counts...) = InfinitePEPS(repeat(T.A, counts...))
 Base.getindex(T::InfinitePEPS, args...) = Base.getindex(T.A, args...)
 Base.setindex!(T::InfinitePEPS, args...) = (Base.setindex!(T.A, args...); T)
 Base.axes(T::InfinitePEPS, args...) = axes(T.A, args...)
+function eachcoordinate(x::InfinitePEPS)
+    return collect(Iterators.product(axes(x)...))
+end
+function eachcoordinate(x::InfinitePEPS, dirs)
+    return collect(Iterators.product(dirs, axes(x, 1), axes(x, 2)))
+end
 TensorKit.space(t::InfinitePEPS, i, j) = space(t[i, j], 1)
 
 ## Math
diff --git a/src/utility/diffable_threads.jl b/src/utility/diffable_threads.jl
new file mode 100644
index 00000000..c03937d9
--- /dev/null
+++ b/src/utility/diffable_threads.jl
@@ -0,0 +1,57 @@
+"""
+    dtmap(args...; kwargs...)
+
+Differentiable wrapper around `OhMyThreads.tmap`.
+
+All calls of `dtmap` inside of PEPSKit use the threading scheduler stored inside
+`Defaults.scheduler` which can be modified using `set_scheduler!`.
+"""
+dtmap(args...; scheduler=Defaults.scheduler[]) = tmap(args...; scheduler)
+
+# Follows the `map` rrule from ChainRules.jl but specified for the case of one AbstractArray that is being mapped
+# https://github.com/JuliaDiff/ChainRules.jl/blob/e245d50a1ae56ce46fc8c1f0fe9b925964f1146e/src/rulesets/Base/base.jl#L243
+function ChainRulesCore.rrule(
+    config::RuleConfig{>:HasReverseMode}, ::typeof(dtmap), f, A::AbstractArray; kwargs...
+)
+    el_rrules = tmap(A; kwargs...) do a
+        rrule_via_ad(config, f, a)
+    end
+    y = map(first, el_rrules)
+    f_projector = ProjectTo(f)
+    A_projectors = map(ProjectTo, A)
+
+    function dtmap_pullback(dy_raw)
+        dys = unthunk(dy_raw)
+        backevals = tmap(el_rrules, dys; kwargs...) do el_rrule, dy
+            last(el_rrule)(dy)
+        end
+        df = f_projector(sum(first, backevals))
+        dA = map((Pa, backeval) -> Pa(last(backeval)), A_projectors, backevals)
+        return NoTangent(), df, dA
+    end
+
+    return y, dtmap_pullback
+end
+
+"""
+    @fwdthreads(ex)
+
+Apply `Threads.@threads` only in the forward pass of the program.
+
+It works by wrapping the for-loop expression in an if statement where in the forward pass
+the loop in computed in parallel using `Threads.@threads`, whereas in the backwards pass
+the `Threads.@threads` is omitted in order to make the expression differentiable.
+"""
+macro fwdthreads(ex)
+    @assert ex.head === :for "@fwdthreads expects a for loop:\n$ex"
+
+    diffable_ex = quote
+        if Zygote.isderiving()
+            $ex
+        else
+            Threads.@threads $ex
+        end
+    end
+
+    return esc(diffable_ex)
+end
diff --git a/src/utility/util.jl b/src/utility/util.jl
index ff881e79..0b9fd4c4 100644
--- a/src/utility/util.jl
+++ b/src/utility/util.jl
@@ -2,6 +2,16 @@
 _next(i, total) = mod1(i + 1, total)
 _prev(i, total) = mod1(i - 1, total)
 
+# iterator over each coordinates
+"""
+    eachcoordinate(x, dirs=1:4)
+
+Enumerate all (dir, row, col) pairs.
+"""
+function eachcoordinate end
+
+@non_differentiable eachcoordinate(args...)
+
 # Element-wise multiplication of TensorMaps respecting block structure
 function _elementwise_mult(a::AbstractTensorMap, b::AbstractTensorMap)
     dst = similar(a)
@@ -166,26 +176,3 @@ macro showtypeofgrad(x)
         end
     )
 end
-
-"""
-    @fwdthreads(ex)
-
-Apply `Threads.@threads` only in the forward pass of the program.
-
-It works by wrapping the for-loop expression in an if statement where in the forward pass
-the loop in computed in parallel using `Threads.@threads`, whereas in the backwards pass
-the `Threads.@threads` is omitted in order to make the expression differentiable.
-"""
-macro fwdthreads(ex)
-    @assert ex.head === :for "@fwdthreads expects a for loop:\n$ex"
-
-    diffable_ex = quote
-        if Zygote.isderiving()
-            $ex
-        else
-            Threads.@threads $ex
-        end
-    end
-
-    return esc(diffable_ex)
-end
diff --git a/test/runtests.jl b/test/runtests.jl
index 2b109130..07a4cd2f 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -20,9 +20,6 @@ end
         @time @safetestset "Unit cell" begin
             include("ctmrg/unitcell.jl")
         end
-        @time @safetestset "SVD wrapper" begin
-            include("ctmrg/svd_wrapper.jl")
-        end
         @time @safetestset ":fixed CTMRG iteration scheme" begin
             include("ctmrg/fixed_iterscheme.jl")
         end
@@ -32,15 +29,23 @@ end
         @time @safetestset "CTMRG schemes" begin
             include("ctmrg/ctmrgschemes.jl")
         end
-        @time @safetestset "CTMRG schemes" begin
-            include("ctmrg/symmetrization.jl")
-        end
     end
     if GROUP == "ALL" || GROUP == "BOUNDARYMPS"
         @time @safetestset "VUMPS" begin
             include("boundarymps/vumps.jl")
         end
     end
+    if GROUP == "ALL" || GROUP == "UTILITY"
+        @time @safetestset "SVD wrapper" begin
+            include("utility/svd_wrapper.jl")
+        end
+        @time @safetestset "Symmetrization" begin
+            include("utility/symmetrization.jl")
+        end
+        @time @safetestset "Differentiable tmap" begin
+            include("utility/diff_maps.jl")
+        end
+    end
     if GROUP == "ALL" || GROUP == "EXAMPLES"
         @time @safetestset "Transverse Field Ising model" begin
             include("tf_ising.jl")
diff --git a/test/utility.jl b/test/test_utils.jl
similarity index 100%
rename from test/utility.jl
rename to test/test_utils.jl
diff --git a/test/utility/diff_maps.jl b/test/utility/diff_maps.jl
new file mode 100644
index 00000000..5606f731
--- /dev/null
+++ b/test/utility/diff_maps.jl
@@ -0,0 +1,7 @@
+using ChainRulesTestUtils
+using PEPSKit: dtmap
+
+# Can the rrule of dtmap be made inferable? (if check_inferred=true, tests error at the moment)
+@testset "Differentiable tmap" begin
+    test_rrule(dtmap, x -> x^3, randn(5, 5); check_inferred=false)
+end
diff --git a/test/ctmrg/svd_wrapper.jl b/test/utility/svd_wrapper.jl
similarity index 100%
rename from test/ctmrg/svd_wrapper.jl
rename to test/utility/svd_wrapper.jl
diff --git a/test/ctmrg/symmetrization.jl b/test/utility/symmetrization.jl
similarity index 100%
rename from test/ctmrg/symmetrization.jl
rename to test/utility/symmetrization.jl