def test_lu_solve_returns_expected_shapes(self, arrays):
m_ss, m_sb, m_bb, m_bs = arrays
hy.assume(hn.nonsquare(m_sb))
hy.assume(hn.all_well_behaved(m_ss))
_, x_f, i_p = gfl.solve_lu(m_ss, m_sb)
expect = utn.array_return_shape('(a,a),(a,b)->(a,b)', m_ss, m_sb)
self.assertArrayShape(gfl.lu_solve(x_f, i_p, m_sb), expect)
expect = utn.array_return_shape('(a,b),(b,b)->(a,b)', m_bs, x_f)
self.assertArrayShape(gfl.rlu_solve(m_bs, x_f, i_p), expect)
with self.assertRaisesRegex(*utn.core_dim_err):
gfl.lu_solve(x_f, i_p, m_bb)
with self.assertRaisesRegex(*utn.core_dim_err):
gfl.rlu_solve(m_sb, x_f, i_p)
_, x_f, i_p = gfl.solve_lu(m_ss, m_sb)
with self.assertRaisesRegex(*utn.broadcast_err):
gfl.lu_solve(x_f, *utn.make_bad_broadcast(i_p, m_sb, (1, 2)))
x_f, i_p = unbroadcast_factors(m_ss, x_f, i_p)
expect = utn.array_return_shape('(a,b),(b,b)->(a,b)', m_bs, m_ss)
self.assertArrayShape(gfl.rlu_solve(m_bs, x_f, i_p), expect)
def test_solve_lu_returns_expected_values(self, arrays):
m_ss, m_sb, m_bs = arrays
hy.assume(hn.all_well_behaved(m_ss))
cond = np.linalg.cond(m_ss).max()
x0_sb = gfl.solve(m_ss, m_sb)
x_sb, x_f, i_p = gfl.solve_lu(m_ss, m_sb)
# with self.subTest('solve0'):
self.assertArrayAllClose(x_sb, x0_sb, cond=cond)
xx_sb = gfl.lu_solve(x_f, i_p, m_sb)
# with self.subTest('solve(lu)'):
self.assertArrayAllClose(xx_sb, x0_sb, cond=cond)
x_bs = gfl.rlu_solve(m_bs, x_f, i_p)
# with self.subTest('rsolve(lu)'):
self.assertArrayAllClose(x_bs @ m_ss, m_bs, cond=cond)
def test_lu_solve_flexible_signature_with_vectors(self, arrays):
m_ss = arrays[0]
v_s, v_b = hn.core_only(*arrays[1:], dims=1)
hy.assume(len(v_s) != len(v_b))
hy.assume(hn.all_well_behaved(m_ss))
# with self.subTest('lu_solve'):
_, x_f, i_p = gfl.solve_lu(m_ss, v_s)
self.assertArrayShape(gfl.lu_solve(x_f, i_p, v_s), m_ss.shape[:-1])
with self.assertRaisesRegex(*utn.core_dim_err):
gfl.lu_solve(x_f, i_p, v_b)
# with self.subTest('rlu_solve'):
self.assertArrayShape(gfl.rlu_solve(v_s, x_f, i_p), m_ss.shape[:-1])
with self.assertRaisesRegex(*utn.core_dim_err):
gfl.rlu_solve(v_b, x_f, i_p)
def test_solve_lu_flexible_signature_with_vectors(self, arrays):
m_ss, m_sb, m_bb, m_bs = arrays[:-1]
v_s = hn.core_only(arrays[-1], dims=1)
hy.assume(hn.nonsquare(m_sb))
hy.assume(hn.all_well_behaved(m_ss))
off_b, y_one = utn.make_off_by_one(m_bb, m_sb)
# with self.subTest('solve_lu'):
self.assertArrayShapesAre(gfl.solve_lu(
m_ss, v_s), (m_ss.shape[:-1], m_ss.shape, m_ss.shape[:-1]))
with self.assertRaisesRegex(*utn.core_dim_err):
gfl.solve_lu(m_bb, v_s)
with self.assertRaisesRegex(*utn.core_dim_err):
gfl.solve_lu(m_bs, v_s)
with self.assertRaisesRegex(*utn.core_dim_err):
# This would succeed/broadcast error if interpreted as Mv:
gfl.solve_lu(m_bb[off_b], m_sb[y_one])
def test_solve_lu_returns_expected_shapes(self, arrays):
m_ss, m_sb, m_bb, m_bs = arrays
hy.assume(hn.nonsquare(m_sb))
hy.assume(hn.all_well_behaved(m_ss))
expect = utn.array_return_shape('(a,a),(a,b)->(a,b)', m_ss, m_sb)
expect_f = expect[:-2] + 2 * m_ss.shape[-1:]
result = gfl.solve_lu(m_ss, m_sb)
self.assertArrayShapesAre(result, (expect, expect_f, expect_f[:-1]))
self.assertArrayShapesAre(unbroadcast_factors(m_ss, *result[1:]),
(m_ss.shape, m_ss.shape[:-1]))
with self.assertRaisesRegex(*utn.core_dim_err):
gfl.solve_lu(m_bb, m_sb)
with self.assertRaisesRegex(*utn.core_dim_err):
gfl.solve_lu(m_bs, m_sb)
with self.assertRaisesRegex(*utn.broadcast_err):
gfl.solve_lu(*utn.make_bad_broadcast(m_ss, m_sb))