def test_set_diag_iterable():
m = Matrix(shape=(4, 4))
m.set_diag(0, range(5))
for i in range(4):
assert m[i, i] == i
def test_set_diag_above():
m = Matrix(shape=(4, 4))
m.set_diag(1, 2)
for i in range(3):
assert m[i, i + 1] == 2.0
def test_set_diag_below():
m = Matrix(shape=(4, 4))
m.set_diag(-1, 2)
for i in range(3):
assert m[i + 1, i] == 2.0
def test_set_diag_float():
m = Matrix(shape=(4, 4))
m.set_diag(0, 2)
for i in range(4):
assert m[i, i] == 2.0
