error retrieving diagonal from rectangular matrix

Greta overloads the R base diag function, but the greta version does not seem to function as advertised; it should extract the main diagonal of a rectangular matrix but instead throws an error. Example from https://greta-stats.org/reference/extract-replace-combine.html below.

``` r
x <- greta::as_data(matrix(1:12, 3, 4))
#> ℹ Initialising python and checking dependencies, this may take a moment.
#> ✔ Initialising python and checking dependencies ... done!
#>
x[1:3, ]
#> greta array (operation)
#>
#> [,1] [,2] [,3] [,4]
#> [1,] 1 4 7 10
#> [2,] 2 5 8 11
#> [3,] 3 6 9 12
x[, 2:4] <- 1:9
e <- greta::diag(x)
#> Error: diagonal elements can only be extracted from square matrices
```