diff --git a/src/math/stdlib_math_diff.fypp b/src/math/stdlib_math_diff.fypp
index e73134607..bad2e7e2c 100644
--- a/src/math/stdlib_math_diff.fypp
+++ b/src/math/stdlib_math_diff.fypp
@@ -75,18 +75,26 @@ contains
         
         size_prepend = 0
         size_append = 0
+        
+        dim_ = 1
         if (present(dim)) then
-            if (dim == 1 .or. dim == 2) then
-                dim_ = dim
-            else
-                dim_ = 1
+            if (dim == 1 .or. dim == 2) dim_ = dim
+        end if
+
+        if (present(prepend)) then
+            if (size(prepend, 3 - dim_) /= size(x, 3 - dim_)) then
+                error stop "stdlib_math_diff: non-differencing dimension of 'prepend' must match 'x'"
+            end if
+            size_prepend = size(prepend, dim_)
+        end if
+        
+        if (present(append)) then
+            if (size(append, 3 - dim_) /= size(x, 3 - dim_)) then
+                error stop "stdlib_math_diff: non-differencing dimension of 'append' must match 'x'"
             end if
-        else
-            dim_ = 1
+            size_append = size(append, dim_)
         end if
         
-        if (present(prepend)) size_prepend = size(prepend, dim_)
-        if (present(append)) size_append = size(append, dim_)
         size_x = size(x, dim_)
         size_work = size_x + size_prepend + size_append
         
@@ -96,8 +104,6 @@ contains
                 allocate(y(0, size(x, 2)))
             case (2)
                 allocate(y(size(x, 1), 0))
-            case default
-                error stop "diff_2: internal error: invalid dimension (dim_ must be 1 or 2)"
             end select
             return
         end if
diff --git a/test/math/test_stdlib_math.fypp b/test/math/test_stdlib_math.fypp
index ea355cd02..72a21a550 100644
--- a/test/math/test_stdlib_math.fypp
+++ b/test/math/test_stdlib_math.fypp
@@ -583,7 +583,7 @@ contains
     end subroutine test_all_close_cmplx_${k1}$
     #:endfor
     
-    #:for k1, t1 in REAL_KINDS_TYPES
+   #:for k1, t1 in REAL_KINDS_TYPES
     subroutine test_diff_real_${k1}$(error)
         type(error_type), allocatable, intent(out) :: error
         ${t1}$ :: x(6) = [${t1}$ :: 0, 5, 15, 30, 50, 75]
@@ -626,6 +626,8 @@ contains
         call check(error, size(diff(B, 2)), 0, "size(diff(B, 2)) in test_diff_real_${k1}$ failed")
         if (allocated(error)) return
         call check(error, size(diff(B, 3)), 0, "size(diff(B, 3)) in test_diff_real_${k1}$ failed")
+        if (allocated(error)) return
+
         
     end subroutine test_diff_real_${k1}$
     #:endfor