From 8caf51b8b7526e97c858bb9e22da49c5bebb7d90 Mon Sep 17 00:00:00 2001
From: Jishnu Bhattacharya <jishnub.github@gmail.com>
Date: Tue, 11 Jul 2023 15:11:39 +0530
Subject: [PATCH] Preserve axes in Diagonal

---
 stdlib/LinearAlgebra/src/diagonal.jl  | 5 +++++
 stdlib/LinearAlgebra/test/diagonal.jl | 6 ++++++
 2 files changed, 11 insertions(+)

diff --git a/stdlib/LinearAlgebra/src/diagonal.jl b/stdlib/LinearAlgebra/src/diagonal.jl
index 37359fd1074f8..3730ac3534af0 100644
--- a/stdlib/LinearAlgebra/src/diagonal.jl
+++ b/stdlib/LinearAlgebra/src/diagonal.jl
@@ -150,6 +150,11 @@ end
     r
 end
 
+function axes(D::Diagonal)
+    ax = axes(D.diag,1)
+    (ax, ax)
+end
+
 @inline function Base.isstored(D::Diagonal, i::Int, j::Int)
     @boundscheck checkbounds(D, i, j)
     if i == j
diff --git a/stdlib/LinearAlgebra/test/diagonal.jl b/stdlib/LinearAlgebra/test/diagonal.jl
index 61045a957cfed..f4c67e901282c 100644
--- a/stdlib/LinearAlgebra/test/diagonal.jl
+++ b/stdlib/LinearAlgebra/test/diagonal.jl
@@ -460,6 +460,12 @@ Random.seed!(1)
     end
 end
 
+@testset "axes" begin
+    v = OffsetArray(1:3)
+    D = Diagonal(v)
+    @test axes(D) isa NTuple{2,typeof(axes(v,1))}
+end
+
 @testset "rdiv! (#40887)" begin
     @test rdiv!(Matrix(Diagonal([2.0, 3.0])), Diagonal(2:3)) == Diagonal([1.0, 1.0])
     @test rdiv!(fill(3.0, 3, 3), 3.0I(3)) == ones(3,3)