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), )