def test_rlstsq_qr_returns_expected_shape_wide(self, arrays, fun):
_, m_sb, m_bb, m_bs = arrays
hy.assume(hn.wide(m_sb))
hy.assume(hn.all_well_behaved(m_sb))
wide = m_sb.shape
expect = utn.array_return_shape('(m,n),(p,n)->(m,p),(n,p)', m_bb, m_sb)
tau = expect[1][:-2] + tau_len(m_sb, fun)
result = fun(m_bb, m_sb)
self.assertArrayShapesAre(result, expect + (tau, ))
self.assertArrayShapesAre(unbroadcast_factors(m_sb, *result[1:]),
(utn.trnsp(wide), wide[:-2] + tau[-1:]))
with self.assertRaisesRegex(*utn.core_dim_err):
fun(m_bs, m_sb)
with self.assertRaisesRegex(*utn.broadcast_err):
fun(*utn.make_bad_broadcast(m_bb, la.transpose(m_bs)))
def test_lstsq_qr_returns_expected_shape_tall(self, arrays, fun):
m_sb, m_bb, m_bs = arrays
hy.assume(hn.tall(m_bs))
hy.assume(hn.all_well_behaved(m_bs))
tall = m_bs.shape
expect = utn.array_return_shape('(m,n),(m,p)->(n,p),(n,m)', m_bs, m_bb)
tau = expect[1][:-2] + tau_len(m_bs, fun)
result = fun(m_bs, m_bb)
self.assertArrayShapesAre(result, expect + (tau, ))
self.assertArrayShapesAre(unbroadcast_factors(m_bs, *result[1:]),
(utn.trnsp(tall), tall[:-2] + tau[-1:]))
with self.assertRaisesRegex(*utn.core_dim_err):
fun(m_bs, m_sb)
with self.assertRaisesRegex(*utn.broadcast_err):
fun(*utn.make_bad_broadcast(m_bs, m_bb))
def test_rsolve_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_bb))
expect = utn.array_return_shape('(a,b),(b,b)->(a,b)', m_sb, m_bb)
expect_f = expect[:-2] + m_bb.shape[-2:]
result = gfl.rsolve_lu(m_sb, m_bb)
self.assertArrayShapesAre(result, (expect, expect_f, expect_f[:-1]))
self.assertArrayShapesAre(unbroadcast_factors(m_bb, *result[1:]),
(m_bb.shape, m_bb.shape[:-1]))
with self.assertRaisesRegex(*utn.core_dim_err):
gfl.rsolve_lu(m_bs, m_bb)
with self.assertRaisesRegex(*utn.core_dim_err):
gfl.rsolve_lu(m_bs, m_sb)
with self.assertRaisesRegex(*utn.broadcast_err):
gfl.rsolve_lu(*utn.make_bad_broadcast(m_bs, m_ss))
def test_rqr_lstsq_returns_expected_shape_wide(self, arrays, fun):
m_ss, m_sb, m_bb, m_bs = arrays
hy.assume(hn.wide(m_sb))
hy.assume(hn.all_well_behaved(m_sb))
_, x_f, tau = fun(m_bb, m_sb)
expect = utn.array_return_shape('(n,m),(m,p)->(n,p)', x_f, m_ss)
self.assertArrayShape(gfl.qr_lstsq(x_f, tau, m_ss), expect)
expect = utn.array_return_shape('(m,n),(n,p)->(m,p)', m_bb, x_f)
self.assertArrayShape(gfl.rqr_lstsq(m_bb, x_f, tau), expect)
with self.assertRaisesRegex(*utn.core_dim_err):
gfl.qr_lstsq(x_f, tau, m_bs)
with self.assertRaisesRegex(*utn.core_dim_err):
gfl.rqr_lstsq(m_bs, x_f, tau)
with self.assertRaisesRegex(*utn.broadcast_err):
gfl.qr_lstsq(x_f, *utn.make_bad_broadcast(tau, m_ss, (1, 2)))
x_f, tau = unbroadcast_factors(m_sb, x_f, tau)
expect = utn.array_return_shape('(m,n),(m,p)->(n,p)', m_sb, m_ss)
self.assertArrayShape(gfl.qr_lstsq(x_f, tau, m_ss), expect)
def test_rlu_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_bb))
_, x_f, i_p = gfl.rsolve_lu(m_sb, m_bb)
expect = utn.array_return_shape('(a,b),(b,b)->(a,b)', m_sb, x_f)
self.assertArrayShape(gfl.rlu_solve(m_sb, x_f, i_p), expect)
expect = utn.array_return_shape('(a,a),(a,b)->(a,b)', x_f, m_bs)
self.assertArrayShape(gfl.lu_solve(x_f, i_p, m_bs), expect)
with self.assertRaisesRegex(*utn.core_dim_err):
gfl.rlu_solve(m_ss, x_f, i_p)
with self.assertRaisesRegex(*utn.core_dim_err):
gfl.lu_solve(x_f, i_p, m_sb)
_, x_f, i_p = gfl.rsolve_lu(m_sb, m_bb)
with self.assertRaisesRegex(*utn.broadcast_err):
gfl.rlu_solve(*utn.make_bad_broadcast(m_sb, x_f), i_p)
x_f, i_p = unbroadcast_factors(m_bb, x_f, i_p)
expect = utn.array_return_shape('(a,a),(a,b)->(a,b)', m_bb, m_bs)
self.assertArrayShape(gfl.lu_solve(x_f, i_p, m_bs), expect)
def test_rqr_lstsq_flexible_signature_with_vectors_mv(self, arrays, fun):
m_sb, m_bs = hn.core_only(*arrays[:-1])
v_s = hn.core_only(arrays[-1], dims=1)
wide, tall = [arr.shape[-2:] for arr in arrays[:-1]]
hy.assume(hn.wide(m_sb))
hy.assume(la.norm(v_s) > 0.)
_, x_f, tau = fun(m_bs, v_s)
expect = utn.array_return_shape('(m,n),(n,p)->(m,p)', m_bs, m_sb)[:-1]
self.assertArrayShape(gfl.qr_lstsq(x_f, tau, m_sb), expect)
self.assertArrayShape(gfl.rqr_lstsq(m_bs, x_f, tau), tall[:-1])
with self.assertRaisesRegex(*utn.core_dim_err):
gfl.qr_lstsq(x_f, tau, m_bs)
self.assertArrayShape(gfl.rqr_lstsq(v_s, x_f, tau), tall[:-2])
m_sb, m_bs = arrays[:-1]
wide, tall = [arr.shape for arr in arrays[:-1]]
_, x_f, tau = fun(m_bs, v_s)
x_f, tau = unbroadcast_factors(v_s, x_f, tau)
self.assertArrayShape(gfl.qr_lstsq(x_f, tau, m_sb), utn.drop(wide))
self.assertArrayShape(gfl.rqr_lstsq(m_bs, x_f, tau), tall[:-1])