def test_matrix(self):
x = matrix()
np_x = np.arange(77).reshape(7, 11).astype(aesara.config.floatX)
fn = aesara.function([x], GpuExtractDiag()(x), mode=mode_with_gpu)
assert np.allclose(fn(np_x), np_x.diagonal())
fn = aesara.function([x], GpuExtractDiag(2)(x), mode=mode_with_gpu)
assert np.allclose(fn(np_x), np_x.diagonal(2))
fn = aesara.function([x], GpuExtractDiag(-3)(x), mode=mode_with_gpu)
assert np.allclose(fn(np_x), np_x.diagonal(-3))
def test_tensor_float16(self):
x = tensor4()
np_x = np.arange(30107).reshape(7, 11, 17, 23).astype("float16")
for offset, axis1, axis2 in [
(1, 0, 1),
(-1, 0, 1),
(0, 1, 0),
(-2, 1, 0),
(-3, 1, 0),
(-2, 2, 0),
(3, 3, 0),
(-1, 3, 2),
(2, 2, 3),
(-1, 2, 1),
(1, 3, 1),
(-1, 1, 3),
]:
assert np.allclose(
GpuExtractDiag(offset, axis1, axis2)(x).eval({x: np_x}),
np_x.diagonal(offset, axis1, axis2),
)