def test_lstsq_flexible_signature_with_vectors_vm(self, arrays):
m_sb, m_bs = arrays[:-2]
v_s, v_b = hn.core_only(*arrays[-2:], dims=1)
wide, tall = [arr.shape for arr in arrays[:-2]]
hy.assume(hn.nonsquare(m_sb))
self.assertArrayShape(gfl.lstsq(v_s, m_sb), utn.drop(wide))
self.assertArrayShape(gfl.lstsq(v_b, m_bs), utn.drop(tall))
with self.assertRaisesRegex(*utn.core_dim_err):
gfl.lstsq(v_s, m_bs)
def test_lstsq_flexible_signature_with_vectors_mv(self, arrays):
m_sb, m_bs = arrays[:-2]
v_s, v_b = hn.core_only(*arrays[-2:], dims=1)
wide, tall = [arr.shape for arr in arrays[:-2]]
hy.assume(hn.wide(m_sb))
off_b, y_one = utn.make_off_by_one(m_sb, m_bs)
self.assertArrayShape(gfl.lstsq(m_bs, v_b), utn.drop(tall))
self.assertArrayShape(gfl.lstsq(m_sb, v_s), utn.drop(wide))
with self.assertRaisesRegex(*utn.core_dim_err):
gfl.lstsq(m_sb, v_b)
with self.assertRaisesRegex(*utn.core_dim_err):
# This would succeed/broadcast error if interpreted as Mv:
gfl.lstsq(m_sb[off_b], m_bs[y_one])
def test_functions_matmul(self, arrays):
m_sb, m_bs = arrays[:-2]
v_s, v_b = hn.core_only(*arrays[-2:], dims=1)
wide, tall = [arr.shape for arr in arrays[:-2]]
# with self.subTest('matmul'):
expect = gf.return_shape('(a,b),(b,c)->(a,c)', tall, wide)
self.assertArrayShape(la.matmul(m_bs, m_sb), expect)
self.assertArrayShape(la.matmul(m_bs, v_s), tall[:-1])
self.assertArrayShape(la.matmul(v_b, m_bs), utn.drop(tall))
self.assertArrayShape(la.matmul(v_s, v_s), ())
# with self.subTest('rmatmul'):
expect = gf.return_shape('(a,b),(b,c)->(a,c)', wide, tall)
self.assertArrayShape(gf.rmatmul(m_bs, m_sb), expect)
self.assertArrayShape(gf.rmatmul(m_sb, v_s), utn.drop(wide))
self.assertArrayShape(gf.rmatmul(v_b, m_sb), wide[:-1])
self.assertArrayShape(gf.rmatmul(v_b, v_b), ())
def test_norm_returns_expected_shapes(self, m_bs):
v_s = m_bs[(0,) * (m_bs.ndim - 1)]
tall = m_bs.shape
self.assertArrayShape(self.gfm.norm(m_bs), tall[:-1])
self.assertArrayShape(self.gfm.norm(m_bs, axis=-2), utn.drop(tall))
self.assertArrayShape(self.gfm.norm(v_s, keepdims=True), (1,))
self.assertArrayShape(self.gfm.norm(v_s), ())
def test_lu(self, arrays):
m_ss, m_sb, m_bs = arrays
smol, wide, tall = [arr.shape for arr in arrays]
hy.assume(hn.wide(m_sb))
# with self.subTest("separate"):
self.assertArrayShapesAre(la.lu(m_ss, 'separate'),
(smol, smol, smol[:-1]))
self.assertArrayShapesAre(la.lu(m_bs, 'separate'),
(tall, utn.chop(tall), utn.drop(tall)))
self.assertArrayShapesAre(la.lu(m_sb, 'separate'),
(utn.chop(wide), wide, wide[:-1]))
# with self.subTest("raw"):
self.assertArrayShapesAre(la.lu(m_ss, 'raw'), (smol, smol[:-1]))
self.assertArrayShapesAre(la.lu(m_bs, 'raw'),
(utn.trnsp(tall), utn.drop(tall)))
self.assertArrayShapesAre(la.lu(m_sb, 'raw'),
(utn.trnsp(wide), wide[:-1]))
def test_lu_basic_returns_expected_shapes(self, arrays):
m_sb, m_bb, m_bs = arrays
wide, big, tall = [arr.shape for arr in arrays]
hy.assume(hn.wide(m_sb))
# with self.subTest(msg="square"):
self.assertArrayShapesAre(gfl.lu_m(m_bb), (big, big, big[:-1]))
self.assertArrayShapesAre(gfl.lu_n(m_bb), (big, big, big[:-1]))
# with self.subTest(msg="wide"):
self.assertArrayShapesAre(gfl.lu_m(m_sb),
(utn.chop(wide), wide, wide[:-1]))
self.assertArrayShapesAre(gfl.lu_n(m_sb),
(wide, utn.grow(wide), utn.drop(wide)))
# with self.subTest(msg="tall"):
self.assertArrayShapesAre(gfl.lu_m(m_bs),
(utn.grow(tall), tall, tall[:-1]))
self.assertArrayShapesAre(gfl.lu_n(m_bs),
(tall, utn.chop(tall), utn.drop(tall)))
def test_lstsq_qr_flexible_signature_with_vectors_mv(self, arrays, fun):
m_sb, m_bs = arrays[:-2]
wide, tall = [arr.shape for arr in arrays[:-2]]
v_s, v_b = hn.core_only(*arrays[-2:], dims=1)
hy.assume(hn.wide(m_sb))
hy.assume(hn.all_well_behaved(m_bs, m_sb))
off_b, y_one = utn.make_off_by_one(m_sb, m_bs)
tau = m_bs.shape[:-2] + tau_len(m_bs, fun)
self.assertArrayShapesAre(fun(m_bs, v_b),
(utn.drop(tall), utn.trnsp(tall), tau))
tau = m_sb.shape[:-2] + tau_len(m_sb, fun)
self.assertArrayShapesAre(fun(m_sb, v_s),
(utn.drop(wide), utn.trnsp(wide), tau))
with self.assertRaisesRegex(*utn.core_dim_err):
fun(m_sb, v_b)
with self.assertRaisesRegex(*utn.core_dim_err):
# This would succeed/broadcast error if interpreted as Mv:
fun(m_sb[off_b], m_bs[y_one])
def test_rqr_lstsq_flexible_signature_with_vectors_vm(self, arrays, fun):
m_sb, m_bs = arrays[:-2]
v_s, v_b = hn.core_only(*arrays[-2:], dims=1)
wide, tall = [arr.shape for arr in arrays[:-2]]
hy.assume(hn.wide(m_sb))
hy.assume(hn.all_well_behaved(m_bs, m_sb))
off_b, y_one = utn.make_off_by_one(m_sb, m_bs)
_, x_f, tau = fun(v_s, m_bs)
self.assertArrayShape(gfl.qr_lstsq(x_f, tau, v_b), utn.drop(tall))
self.assertArrayShape(gfl.rqr_lstsq(v_s, x_f, tau), tall[:-1])
_, x_f, tau = fun(v_b, m_sb)
self.assertArrayShape(gfl.qr_lstsq(x_f, tau, v_s), utn.drop(wide))
self.assertArrayShape(gfl.rqr_lstsq(v_b, x_f, tau), wide[:-1])
with self.assertRaisesRegex(*utn.core_dim_err):
gfl.rqr_lstsq(v_s, x_f, tau)
with self.assertRaisesRegex(*utn.core_dim_err):
# This would succeed/broadcast error if interpreted as Mv:
gfl.rqr_lstsq(m_bs[y_one], x_f[off_b], tau[off_b])
def test_rlstsq_qr_flexible_signature_with_vectors_mv(self, arrays, fun):
m_sb, m_bs = arrays[:-1]
wide = m_sb.shape
v_b = hn.core_only(arrays[-1], dims=1)
hy.assume(hn.wide(m_sb))
hy.assume(la.norm(v_b) > 0.)
tau = m_sb.shape[:-2] + tau_len_vec(v_b, fun)
self.assertArrayShapesAre(fun(m_sb, v_b),
(wide[:-1], utn.drop(wide), tau))
with self.assertRaisesRegex(*utn.core_dim_err):
fun(m_bs, v_b)
def test_qr_lstsq_flexible_signature_with_vectors_vm(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(v_s, m_sb)
self.assertArrayShape(gfl.qr_lstsq(x_f, tau, m_sb), utn.drop(wide))
expect = utn.array_return_shape('(m,n),(n,p)->(m,p)', m_bs, m_sb)[:-1]
self.assertArrayShape(gfl.rqr_lstsq(m_bs, x_f, tau), expect)
with self.assertRaisesRegex(*utn.core_dim_err):
gfl.qr_lstsq(x_f, tau, m_bs)
self.assertArrayShape(gfl.qr_lstsq(x_f, tau, v_s), wide[:-2])
# with broadcasting
m_sb, m_bs = arrays[:-1]
wide, tall = [arr.shape for arr in arrays[:-1]]
_, x_f, tau = fun(v_s, m_sb)
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])
def test_functions_matldiv(self, arrays):
m_ss, m_sb, m_bb, m_bs = arrays[:-1]
v_s = hn.core_only(arrays[-1], dims=1)
smol, wide, big, tall = [arr.shape for arr in arrays[:-1]]
hy.assume(hn.all_well_behaved(m_ss))
hy.assume(hn.wide(m_sb))
# with self.subTest('solve'):
expect = gf.return_shape('(a,b),(a,c)->(b,c)', smol, wide)
self.assertArrayShape(la.matldiv(m_ss, m_sb), expect)
self.assertArrayShape(la.matldiv(m_ss, v_s), smol[:-1])
# with self.subTest('lstsq'):
expect = gf.return_shape('(a,b),(a,c)->(b,c)', tall, big)
self.assertArrayShape(la.matldiv(m_bs, m_bb), expect)
expect = gf.return_shape('(a,b),(a,c)->(b,c)', wide, smol)
self.assertArrayShape(la.matldiv(m_sb, m_ss), expect)
self.assertArrayShape(la.matldiv(m_sb, v_s), utn.drop(wide))
self.assertArrayShape(la.matldiv(v_s, m_ss), smol[:-1])
def test_functions_lstsq(self, arrays):
m_ss, m_sb, m_bb, m_bs = arrays[:-2]
v_s, v_b = hn.core_only(*arrays[-2:], dims=1)
smol, wide, big, tall = [arr.shape for arr in arrays[:-2]]
# with self.subTest('lstsq'):
expect = gf.return_shape('(a,b),(a,c)->(b,c)', tall, big)
self.assertArrayShape(la.lstsq(m_bs, m_bb), expect)
expect = gf.return_shape('(a,b),(a,c)->(b,c)', wide, smol)
self.assertArrayShape(la.lstsq(m_sb, m_ss), expect)
self.assertArrayShape(la.lstsq(m_sb, v_s), utn.drop(wide))
self.assertArrayShape(la.lstsq(v_s, m_ss), smol[:-1])
# with self.subTest('rlstsq'):
expect = gf.return_shape('(a,b),(c,b)->(a,c)', smol, tall)
self.assertArrayShape(la.rlstsq(m_ss, m_bs), expect)
expect = gf.return_shape('(a,b),(c,b)->(a,c)', big, wide)
self.assertArrayShape(la.rlstsq(m_bb, m_sb), expect)
self.assertArrayShape(la.rlstsq(v_b, m_sb), wide[:-1])
self.assertArrayShape(la.rlstsq(m_ss, v_s), smol[:-1])
self.assertArrayShape(la.rlstsq(v_s, m_bs), tall[:-1])
def test_lq_returns_expected_shapes(self, arrays):
m_sb, m_bs = arrays
wide, tall = [arr.shape for arr in arrays]
hy.assume(hn.wide(m_sb))
# with self.subTest(msg='wide'):
self.assertArrayShapesAre(gfl.lq_m(m_sb), (utn.chop(wide), wide))
self.assertArrayShape(gfl.lq_lm(m_sb), utn.chop(wide))
# with self.subTest(msg='tall'):
self.assertArrayShapesAre(gfl.lq_n(m_bs), (tall, utn.chop(tall)))
self.assertArrayShape(gfl.lq_ln(m_bs), tall)
with self.assertRaisesRegex(*utn.invalid_err):
gfl.lq_m(m_bs)
# with self.subTest(msg='complete'):
self.assertArrayShapesAre(gfl.lq_n(m_sb), (wide, utn.grow(wide)))
# with self.subTest(msg='raw'):
self.assertArrayShapesAre(gfl.lq_rawm(m_sb),
(utn.trnsp(wide), wide[:-1]))
self.assertArrayShapesAre(gfl.lq_rawn(m_bs),
(utn.trnsp(tall), utn.drop(tall)))
def test_lq(self, arrays):
m_sb, m_bs = arrays
wide, tall = [arr.shape for arr in arrays]
hy.assume(hn.wide(m_sb))
# with self.subTest("reduced"):
self.assertArrayShapesAre(la.lq(m_bs, 'reduced'),
(tall, utn.chop(tall)))
self.assertArrayShapesAre(la.lq(m_sb, 'reduced'),
(utn.chop(wide), wide))
# with self.subTest("complete"):
self.assertArrayShapesAre(la.lq(m_bs, 'complete'),
(tall, utn.chop(tall)))
self.assertArrayShapesAre(la.lq(m_sb, 'complete'),
(wide, utn.grow(wide)))
# with self.subTest("l/raw"):
self.assertArrayShape(la.lq(m_bs, 'l'), tall)
self.assertArrayShape(la.lq(m_sb, 'l'), utn.chop(wide))
self.assertArrayShapesAre(la.lq(m_bs, 'raw'),
(utn.trnsp(tall), utn.drop(tall)))
self.assertArrayShapesAre(la.lq(m_sb, 'raw'),
(utn.trnsp(wide), wide[:-1]))
def test_pinv_returns_expected_shapes(self, arrays):
m_sb, m_bs = arrays
wide, tall = [arr.shape for arr in arrays]
hy.assume(hn.wide(m_sb))
hy.assume(hn.all_well_behaved(m_bs, m_sb))
# with self.subTest(msg='wide'):
self.assertArrayShape(gfl.pinv(m_sb), utn.trnsp(wide))
# with self.subTest(msg='tall'):
self.assertArrayShape(gfl.pinv(m_bs), utn.trnsp(tall))
# with self.subTest(msg='wide,+qr'):
self.assertArrayShapesAre(
gfl.pinv_qrm(m_sb), (utn.trnsp(wide), utn.trnsp(wide), wide[:-1]))
# with self.subTest(msg='tall,+qr'):
self.assertArrayShapesAre(
gfl.pinv_qrn(m_bs),
(utn.trnsp(tall), utn.trnsp(tall), utn.drop(tall)))
# with self.subTest(msg='wide,-qr'):
_, m_sb_f, m_sb_tau = gfl.pinv_qrm(m_sb)
self.assertArrayShape(gfl.qr_pinv(m_sb_f, m_sb_tau), utn.trnsp(wide))
# with self.subTest(msg='tall,-qr'):
_, m_bs_f, m_bs_tau = gfl.pinv_qrn(m_bs)
self.assertArrayShape(gfl.qr_pinv(m_bs_f, m_bs_tau), utn.trnsp(tall))
