diff --git a/Project.toml b/Project.toml
index e870a7a4..2678d52c 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "DynamicalSystemsBase"
 uuid = "6e36e845-645a-534a-86f2-f5d4aa5a06b4"
 repo = "https://github.com/JuliaDynamics/DynamicalSystemsBase.jl.git"
-version = "3.11.0"
+version = "3.11.1"
 
 [deps]
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
diff --git a/ext/src/CoupledSDEs.jl b/ext/src/CoupledSDEs.jl
index a8e73964..31e37d35 100644
--- a/ext/src/CoupledSDEs.jl
+++ b/ext/src/CoupledSDEs.jl
@@ -195,10 +195,11 @@ If this is not the case, returns `nothing`.
 Note: The diffusion matrix ``Σ`` is the square root of the noise covariance matrix ``Q`` (see
 [`covariance_matrix`](@ref)), defined via the Cholesky decomposition ``Q = Σ Σ^\\top``.
 """
-function diffusion_matrix(ds::CoupledSDEs{IIP,D}) where {IIP,D}
+function diffusion_matrix(ds::CoupledSDEs{IIP,D})::AbstractMatrix where {IIP,D}
     if ds.noise_type[:invertible]
         diffusion = diffusion_function(ds)
         A = diffusion(zeros(D), current_parameters(ds), 0.0)
+        A = A isa AbstractMatrix ? A : Diagonal(A)
     else
         @warn """
         The diffusion function of the `CoupledSDEs` cannot be expressed as a constant
@@ -218,7 +219,7 @@ If this is not the case, returns `nothing`.
 
 See also [`diffusion_matrix`](@ref).
 """
-function covariance_matrix(ds::CoupledSDEs)
+function covariance_matrix(ds::CoupledSDEs)::AbstractMatrix
     A = diffusion_matrix(ds)
     (A == nothing) ? nothing : A * A'
 end
@@ -241,10 +242,10 @@ Returns `g, noise_prototype`.
 """
 function construct_diffusion_function(
     g, covariance, noise_prototype, noise_strength, D, IIP
-    )
+)
     if isnothing(g) # diagonal additive noise
         cov = isnothing(covariance) ? LinearAlgebra.I(D) : covariance
-        size(cov) != (D,D) &&
+        size(cov) != (D, D) &&
             throw(ArgumentError("Covariance matrix must be of size $((D, D))"))
         A = sqrt(cov)
         if IIP
@@ -269,4 +270,4 @@ function construct_diffusion_function(
         end
     end
     return g, noise_prototype
-end
\ No newline at end of file
+end
diff --git a/test/stochastic.jl b/test/stochastic.jl
index 84cacbf1..f783eb99 100644
--- a/test/stochastic.jl
+++ b/test/stochastic.jl
@@ -127,6 +127,7 @@ end
         A = sqrt(Γ)
         lorenz_oop = CoupledSDEs(lorenz_rule, u0, p0, covariance = Γ, diffeq=diffeq_cov)
         @test A ≈ diffusion_matrix(lorenz_oop)
+        @test A isa AbstractMatrix
         @test Γ ≈ A * A'
     end
 
@@ -153,4 +154,4 @@ end
         @test approx ≈ Γ atol=1e-1 broken = true
         # I think I understand something wromg here
     end
-end
\ No newline at end of file
+end