diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
index 01b5a56..2194b51 100644
--- a/.github/workflows/CI.yml
+++ b/.github/workflows/CI.yml
@@ -31,7 +31,6 @@ jobs:
       - uses: julia-actions/setup-julia@v3
         with:
           version: ${{ matrix.version }}
-          arch: x64
       - uses: julia-actions/cache@v3
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
@@ -47,7 +46,7 @@ jobs:
     name: Runic formatting
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v5
+      - uses: actions/checkout@v7
       - uses: fredrikekre/runic-action@v1
         with:
           version: '1'
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
new file mode 100644
index 0000000..94e37a7
--- /dev/null
+++ b/.pre-commit-config.yaml
@@ -0,0 +1,27 @@
+repos:
+  - repo: https://github.com/fredrikekre/runic-pre-commit
+    rev: v1.0.0
+    hooks:
+      - id: runic
+
+  - repo: https://github.com/crate-ci/typos
+    rev: v1.48.0
+    hooks:
+      - id: typos
+        # report and fail (no --write-changes, so no in-place edits);
+        # --force-exclude makes staged files honor _typos.toml excludes
+        args: [--force-exclude]
+
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v6.0.0
+    hooks:
+      - id: trailing-whitespace
+        # ReferenceTests golden files: trailing space and the missing
+        # final newline are significant output, never "fix" them
+        exclude: ^test/references/
+      - id: end-of-file-fixer
+        exclude: ^test/references/
+      - id: check-merge-conflict
+      - id: check-added-large-files
+      - id: check-toml
+      - id: check-yaml
diff --git a/Project.toml b/Project.toml
index 3e7ee5b..856bee7 100644
--- a/Project.toml
+++ b/Project.toml
@@ -9,6 +9,7 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 Metalhead = "dbeba491-748d-5e0e-a39e-b530a07fa0cc"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
@@ -28,6 +29,7 @@ Distributions = "0.25"
 Flux = "0.14, 0.15, 0.16"
 Metalhead = "0.9.5"
 NNlib = "0.9"
+Optimisers = "0.3, 0.4"
 ProgressMeter = "1"
 Random = "1"
 Reexport = "1"
diff --git a/README.md b/README.md
index 824a85c..0abbd5f 100644
--- a/README.md
+++ b/README.md
@@ -1,5 +1,13 @@
 # SmoothDiff.jl
 
+[![Code Style: Runic](https://img.shields.io/badge/code_style-%E1%9A%B1%E1%9A%A2%E1%9A%BE%E1%9B%81%E1%9A%B2-black)](https://github.com/fredrikekre/Runic.jl)
+[![Aqua](https://raw.githubusercontent.com/JuliaTesting/Aqua.jl/master/badge.svg)](https://github.com/JuliaTesting/Aqua.jl)
+[![JET](https://img.shields.io/badge/%F0%9F%9B%A9%EF%B8%8F_tested_with-JET.jl-233f9a)](https://github.com/aviatesk/JET.jl)
+
 Julia reference implementation of SmoothDiff for the NeurIPS 2025 paper *"Smoothed Differentiation Efficiently Mitigates Shattered Gradients in Explanations"*.
 
-The full experiments can be found here: https://github.com/adrhill/smoothdiff-experiments/
\ No newline at end of file
+The full experiments can be found here: https://github.com/adrhill/smoothdiff-experiments/
+
+## Acknowledgements
+
+Adrian Hill gratefully acknowledges funding from the German Federal Ministry of Education and Research under the grant BIFOLD26B.
diff --git a/src/SmoothedDifferentiation.jl b/src/SmoothedDifferentiation.jl
index ef2d282..804660d 100644
--- a/src/SmoothedDifferentiation.jl
+++ b/src/SmoothedDifferentiation.jl
@@ -5,15 +5,16 @@ using Reexport
 import XAIBase: call_analyzer
 
 using Base.Iterators
-using Distributions: Distribution, Sampleable, Normal
+using Distributions: Sampleable, Normal
 using Random: AbstractRNG, GLOBAL_RNG, rand!
 using ProgressMeter: Progress, next!
 
-using NNlib: relu, ∇maxpool, maxpool, upsample_nearest, σ, softplus
+using NNlib: relu, ∇maxpool, maxpool, upsample_nearest, σ
 using Zygote: pullback
 import ChainRulesCore: rrule, NoTangent, unthunk
 
 using Flux: Flux
+using Optimisers: Optimisers
 
 include("prepare_model.jl")
 include("vejp/relu.jl")
@@ -27,7 +28,7 @@ mytestmode!(model) = model
 mytestmode!(model::Flux.Chain) = Flux.testmode!(model)
 
 
-samplingmode!(model, mode::Bool) = foreach(x -> samplingmode!(x, mode), Flux.trainable(model))
+samplingmode!(model, mode::Bool) = foreach(x -> samplingmode!(x, mode), Optimisers.trainable(model))
 
 ## Interface
 const DEFAULT_SAMPLES = 50
diff --git a/test/Project.toml b/test/Project.toml
index e2fcafa..758d2ad 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -1,5 +1,6 @@
 [deps]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
diff --git a/test/linting.jl b/test/linting.jl
new file mode 100644
index 0000000..6fb49f0
--- /dev/null
+++ b/test/linting.jl
@@ -0,0 +1,30 @@
+# Shared code-quality (Aqua.jl), type-stability (JET.jl) and import-hygiene
+# (ExplicitImports.jl) checks.
+# Kept consistent across the Julia-XAI packages — see REFACTOR.md.
+using SmoothedDifferentiation
+using Test
+using Aqua
+using JET
+using ExplicitImports
+
+@testset "Aqua.jl" begin
+    @info "Running Aqua.jl code-quality tests. These might print warnings from dependencies."
+    Aqua.test_all(SmoothedDifferentiation; ambiguities = false)
+    Aqua.test_ambiguities(SmoothedDifferentiation)
+end
+
+@testset "ExplicitImports.jl" begin
+    @info "Running ExplicitImports.jl import-hygiene tests."
+    @test check_no_stale_explicit_imports(SmoothedDifferentiation) === nothing
+    @test check_all_explicit_imports_via_owners(SmoothedDifferentiation) === nothing
+    @test check_all_qualified_accesses_via_owners(SmoothedDifferentiation) === nothing
+    @test check_no_self_qualified_accesses(SmoothedDifferentiation) === nothing
+end
+
+# JET's v0.11 series supports Julia v1.12 and above only.
+if VERSION >= v"1.12"
+    @testset "JET.jl" begin
+        @info "Running JET.jl type-stability tests."
+        JET.test_package(SmoothedDifferentiation; target_defined_modules = true)
+    end
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index 585a0d8..c852ce6 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,14 +1,10 @@
 using SmoothedDifferentiation
 using Test
-using Aqua
-using JET
 
 @testset "SmoothedDifferentiation.jl" begin
-    @testset "Code quality (Aqua.jl)" begin
-        Aqua.test_all(SmoothedDifferentiation)
-    end
-    @testset "Code linting (JET.jl)" begin
-        JET.test_package(SmoothedDifferentiation; target_defined_modules = true)
+    @testset verbose = true "Linting" begin
+        @info "Running linting tests..."
+        include("linting.jl")
     end
 
     @testset "VEJP numerical tests" begin