diff options
Diffstat (limited to 'numpy/lib/twodim_base.py')
| -rw-r--r-- | numpy/lib/twodim_base.py | 8 |
1 files changed, 7 insertions, 1 deletions
diff --git a/numpy/lib/twodim_base.py b/numpy/lib/twodim_base.py index 58d8250a1..2b518aeae 100644 --- a/numpy/lib/twodim_base.py +++ b/numpy/lib/twodim_base.py @@ -210,7 +210,13 @@ def eye(N, M=None, k=0, dtype=float, maskna=False): if M is None: M = N m = zeros((N, M), dtype=dtype, maskna=maskna) - diagonal(m, k)[...] = 1 + if k >= M: + return m + if k >= 0: + i = k + else: + i = (-k) * M + m[:M-k].flat[i::M+1] = 1 return m def diag(v, k=0): |