How to define a DIAG function?

Question

I would like to define a VBA function that:

1. returns a column vector of the main diagonal elements of a matrix;
2. returns a square diagonal matrix with the elements of vector on the main diagonal;
3. automatically returns a matrix/vector, without the need to press Ctrl Shift Enter;

I am working on this code:

Public Function DIAG(matrix As Variant) As Variant

Dim i As Long
Dim j As Long
Dim nRows As Long
Dim nCols As Long
Dim tempArray As Variant

nRows = matrix.Rows.Count
nCols = matrix.Columns.Count

For i = 1 To nRows
    For j = 1 To nCols
        If i = j Then
           tempArray(i) = matrix(i, j)
        End If
    Next j
Next i

DIAG = tempArray

End Function

which is only for the first purpose of the function, but it is not working. I get:

#VALUE

Answer 1

The code from @FaneDuru did help me, but I successfully coded my UDF in this way:

Public Function DIAG(matrix As Range) As Variant

Dim i As Long, j As Long, nRows As Long, nCols As Long
Dim tempArray As Variant

nRows = matrix.Rows.Count
nCols = matrix.Columns.Count

If nCols = 1 Then
 
 ReDim tempArray(nRows - 1, nRows - 1)

 For i = 1 To nRows
  tempArray(i - 1, i - 1) = matrix(i)
 Next i

Else

 If nCols = nRows Then

 ReDim tempArray(nRows - 1, 0)

 For i = 1 To nRows
  For j = 1 To nCols
   If i = j Then
    tempArray(i - 1, 0) = matrix(i, j)
   End If
  Next j
 Next i

 Else
  tempArray = CVErr(xlErrValue)
 End If

End If

DIAG = tempArray

End Function