def test_clutrr_v2():
embedding_size = 20
triples, hops = [], []
xxx = []
for i in range(16):
triples += [(f'a{i}', 'p', f'b{i}'), (f'b{i}', 'q', f'c{i}')]
hops += [(f'a{i}', 'r', f'c{i}')]
xxx += [(f'a{i}', 'p', f'c{i}'), (f'a{i}', 'q', f'c{i}'), (f'a{i}', 'r', f'c{i}')]
entity_lst = sorted({s for (s, _, _) in triples + hops} | {o for (_, _, o) in triples + hops})
predicate_lst = sorted({p for (_, p, _) in triples + hops})
nb_entities, nb_predicates = len(entity_lst), len(predicate_lst)
entity_to_index = {e: i for i, e in enumerate(entity_lst)}
predicate_to_index = {p: i for i, p in enumerate(predicate_lst)}
kernel = GaussianKernel()
entity_embeddings = nn.Embedding(nb_entities, embedding_size, sparse=True)
predicate_embeddings = nn.Embedding(nb_predicates, embedding_size, sparse=True)
for scoring_type in ['concat']: # ['min', 'concat']:
model = NeuralKB(kernel=kernel, scoring_type=scoring_type)
indices = torch.LongTensor(np.array([predicate_to_index['p'], predicate_to_index['q']]))
_hops = SymbolicReformulator(predicate_embeddings, indices)
hoppy = Hoppy(model, hops_lst=[(_hops, False)], depth=1)
for s in entity_lst:
for p in predicate_lst:
for o in entity_lst:
xs_np = np.array([entity_to_index[s]])
xp_np = np.array([predicate_to_index[p]])
xo_np = np.array([entity_to_index[o]])
with torch.no_grad():
xs = torch.LongTensor(xs_np)
xp = torch.LongTensor(xp_np)
xo = torch.LongTensor(xo_np)
xs_emb = entity_embeddings(xs)
xp_emb = predicate_embeddings(xp)
xo_emb = entity_embeddings(xo)
rel_emb = encode_relation(facts=triples, relation_embeddings=predicate_embeddings,
relation_to_idx=predicate_to_index)
arg1_emb, arg2_emb = encode_arguments(facts=triples, entity_embeddings=entity_embeddings,
entity_to_idx=entity_to_index)
facts = [rel_emb, arg1_emb, arg2_emb]
inf = hoppy.score(xp_emb, xs_emb, xo_emb, facts=facts,
entity_embeddings=entity_embeddings.weight)
inf_np = inf.cpu().numpy()
print(s, p, o, inf_np)
assert inf_np[0] > 0.9 if (s, p, o) in (triples + xxx) else inf_np[0] < 0.1
def test_masking_v2():
nb_entities = 10
nb_predicates = 5
embedding_size = 10
rs = np.random.RandomState(0)
for _ in range(1):
for position in [0, 1, 2]:
for st in ['min', 'concat']:
with torch.no_grad():
triples = [
('a', 'p', 'b'),
('b', 'q', 'c'),
('a', 'p', 'c')
]
entity_to_index = {'a': 0, 'b': 1, 'c': 2, 'd': 3}
predicate_to_index = {'p': 0, 'q': 1}
kernel = GaussianKernel()
entity_emb = nn.Embedding(nb_entities, embedding_size * 2, sparse=True)
predicate_emb = nn.Embedding(nb_predicates, embedding_size * 2, sparse=True)
fact_rel = torch.LongTensor(np.array([predicate_to_index[p] for (_, p, _) in triples]))
fact_arg1 = torch.LongTensor(np.array([entity_to_index[s] for (s, _, _) in triples]))
fact_arg2 = torch.LongTensor(np.array([entity_to_index[o] for (_, _, o) in triples]))
facts = [fact_rel, fact_arg1, fact_arg2]
base = NeuralKB(entity_embeddings=entity_emb, predicate_embeddings=predicate_emb,
kernel=kernel, facts=facts, scoring_type=st)
indices = torch.LongTensor(np.array([predicate_to_index['p'], predicate_to_index['q']]))
reformulator = SymbolicReformulator(predicate_emb, indices)
model = SimpleHoppy(base, entity_emb, hops=reformulator)
xs_np = rs.randint(nb_entities, size=32)
xp_np = rs.randint(nb_predicates, size=32)
xo_np = rs.randint(nb_entities, size=32)
xi_np = np.array([position] * xs_np.shape[0])
xs_np[0] = 0
xp_np[0] = 0
xo_np[0] = 1
xs_np[1] = 1
xp_np[1] = 1
xo_np[1] = 2
xs_np[2] = 0
xp_np[2] = 0
xo_np[2] = 2
xs = torch.LongTensor(xs_np)
xp = torch.LongTensor(xp_np)
xo = torch.LongTensor(xo_np)
xi = torch.LongTensor(xi_np)
xs_emb = entity_emb(xs)
xp_emb = predicate_emb(xp)
xo_emb = entity_emb(xo)
# xi = None
base.mask_indices = xi
scores = model.forward(xp_emb, xs_emb, xo_emb)
inf = model.score(xp_emb, xs_emb, xo_emb)
if position in {0, 1}:
assert inf[2] < 0.5
else:
assert inf[2] > 0.9
scores_sp, scores_po = scores
inf = inf.cpu().numpy()
scores_sp = scores_sp.cpu().numpy()
scores_po = scores_po.cpu().numpy()
for i in range(xs.shape[0]):
np.testing.assert_allclose(inf[i], scores_sp[i, xo[i]], rtol=1e-5, atol=1e-5)
np.testing.assert_allclose(inf[i], scores_po[i, xs[i]], rtol=1e-5, atol=1e-5)
def test_clutrr_v7():
torch.set_num_threads(multiprocessing.cpu_count())
embedding_size = 50
torch.manual_seed(0)
rs = np.random.RandomState(0)
triples = [('a', 'p', 'b'), ('b', 'q', 'c'), ('c', 'p', 'd'),
('d', 'q', 'e'), ('e', 'p', 'f'), ('f', 'q', 'g'),
('g', 'p', 'h'), ('h', 'q', 'i'), ('i', 'p', 'l'),
('l', 'q', 'm'), ('m', 'p', 'n'), ('n', 'q', 'o'),
('o', 'p', 'p'), ('p', 'q', 'q'), ('q', 'p', 'r'),
('r', 'q', 's'), ('s', 'p', 't'), ('t', 'q', 'u'),
('u', 'p', 'v'), ('v', 'q', 'w'), ('x', 'r', 'y'),
('x', 's', 'y')]
entity_lst = sorted({e
for (e, _, _) in triples}
| {e
for (_, _, e) in triples})
predicate_lst = sorted({p for (_, p, _) in triples})
nb_entities = len(entity_lst)
nb_predicates = len(predicate_lst)
entity_to_index = {e: i for i, e in enumerate(entity_lst)}
predicate_to_index = {p: i for i, p in enumerate(predicate_lst)}
with torch.no_grad():
kernel = GaussianKernel()
entity_embeddings = nn.Embedding(nb_entities,
embedding_size * 2,
sparse=True)
predicate_embeddings = nn.Embedding(nb_predicates,
embedding_size * 2,
sparse=True)
rel_emb = encode_relation(facts=triples,
relation_embeddings=predicate_embeddings,
relation_to_idx=predicate_to_index)
arg1_emb, arg2_emb = encode_arguments(
facts=triples,
entity_embeddings=entity_embeddings,
entity_to_idx=entity_to_index)
facts = [rel_emb, arg1_emb, arg2_emb]
k = 5
model = NeuralKB(kernel=kernel)
indices = torch.from_numpy(
np.array([predicate_to_index['p'], predicate_to_index['q']]))
reformulator = SymbolicReformulator(predicate_embeddings, indices)
hoppy0 = Hoppy(model, hops_lst=[(reformulator, False)], depth=0)
hoppy1 = Hoppy(model, hops_lst=[(reformulator, False)], depth=1)
hoppy2 = Hoppy(model, hops_lst=[(reformulator, False)], depth=2)
hoppy3 = Hoppy(model, hops_lst=[(reformulator, False)], depth=3)
hoppy4 = Hoppy(model, hops_lst=[(reformulator, False)], depth=4)
xs_np = rs.randint(nb_entities, size=12)
xp_np = rs.randint(nb_predicates, size=12)
xo_np = rs.randint(nb_entities, size=12)
xs_np[0] = entity_to_index['a']
xp_np[0] = predicate_to_index['r']
xo_np[0] = entity_to_index['c']
xs_np[1] = entity_to_index['a']
xp_np[1] = predicate_to_index['r']
xo_np[1] = entity_to_index['e']
xs_np[2] = entity_to_index['a']
xp_np[2] = predicate_to_index['r']
xo_np[2] = entity_to_index['g']
xs_np[3] = entity_to_index['a']
xp_np[3] = predicate_to_index['r']
xo_np[3] = entity_to_index['i']
xs_np[4] = entity_to_index['a']
xp_np[4] = predicate_to_index['r']
xo_np[4] = entity_to_index['m']
xs_np[5] = entity_to_index['a']
xp_np[5] = predicate_to_index['r']
xo_np[5] = entity_to_index['o']
xs_np[6] = entity_to_index['a']
xp_np[6] = predicate_to_index['r']
xo_np[6] = entity_to_index['q']
xs_np[7] = entity_to_index['a']
xp_np[7] = predicate_to_index['r']
xo_np[7] = entity_to_index['s']
xs_np[8] = entity_to_index['a']
xp_np[8] = predicate_to_index['r']
xo_np[8] = entity_to_index['u']
xs = torch.from_numpy(xs_np)
xp = torch.from_numpy(xp_np)
xo = torch.from_numpy(xo_np)
xs_emb = entity_embeddings(xs)
xp_emb = predicate_embeddings(xp)
xo_emb = entity_embeddings(xo)
# res0 = hoppy0.forward(xp_emb, xs_emb, xo_emb, facts=facts, entity_embeddings=entity_embeddings)
inf0 = hoppy0.score(xp_emb,
xs_emb,
xo_emb,
facts=facts,
entity_embeddings=entity_embeddings.weight)
# (scores0_sp, subs0_sp), (scores0_po, subs0_po) = res0
# res1 = hoppy1.forward(xp_emb, xs_emb, xo_emb, facts=facts, entity_embeddings=entity_embeddings)
inf1 = hoppy1.score(xp_emb,
xs_emb,
xo_emb,
facts=facts,
entity_embeddings=entity_embeddings.weight)
# (scores1_sp, subs1_sp), (scores1_po, subs1_po) = res1
# res2 = hoppy2.forward(xp_emb, xs_emb, xo_emb, facts=facts, entity_embeddings=entity_embeddings)
inf2 = hoppy2.score(xp_emb,
xs_emb,
xo_emb,
facts=facts,
entity_embeddings=entity_embeddings.weight)
# (scores2_sp, subs2_sp), (scores2_po, subs2_po) = res2
# res3 = hoppy3.forward(xp_emb, xs_emb, xo_emb, facts=facts, entity_embeddings=entity_embeddings)
inf3 = hoppy3.score(xp_emb,
xs_emb,
xo_emb,
facts=facts,
entity_embeddings=entity_embeddings.weight)
# (scores3_sp, subs3_sp), (scores3_po, subs3_po) = res3
# res4 = hoppy4.forward(xp_emb, xs_emb, xo_emb, facts=facts, entity_embeddings=entity_embeddings)
inf4 = hoppy4.score(xp_emb,
xs_emb,
xo_emb,
facts=facts,
entity_embeddings=entity_embeddings.weight)
# (scores4_sp, subs4_sp), (scores4_po, subs4_po) = res4
inf0_np = inf0.cpu().numpy()
inf1_np = inf1.cpu().numpy()
inf2_np = inf2.cpu().numpy()
inf3_np = inf3.cpu().numpy()
inf4_np = inf4.cpu().numpy()
np.testing.assert_allclose(inf0_np,
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
rtol=1e-1,
atol=1e-1)
np.testing.assert_allclose(inf1_np,
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
rtol=1e-1,
atol=1e-1)
np.testing.assert_allclose(inf2_np,
[1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
rtol=1e-1,
atol=1e-1)
np.testing.assert_allclose(inf3_np,
[1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
rtol=1e-1,
atol=1e-1)
np.testing.assert_allclose(inf4_np,
[1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0],
rtol=1e-1,
atol=1e-1)
print(inf3_np)
print(entity_embeddings.weight[entity_to_index['c'], 0].item())
print(entity_embeddings.weight[entity_to_index['e'], 0].item())
print(entity_embeddings.weight[entity_to_index['g'], 0].item())
print(entity_embeddings.weight[entity_to_index['i'], 0].item())
def test_reasoning_v6():
torch.set_num_threads(multiprocessing.cpu_count())
embedding_size = 50
torch.manual_seed(0)
rs = np.random.RandomState(0)
triples = [('a', 'p', 'b'), ('b', 'q', 'c'), ('c', 'p', 'd'),
('d', 'q', 'e'), ('e', 'p', 'f'), ('f', 'q', 'g'),
('g', 'p', 'h'), ('h', 'q', 'i'), ('i', 'p', 'l'),
('l', 'q', 'm'), ('m', 'p', 'n'), ('n', 'q', 'o'),
('o', 'p', 'p'), ('p', 'q', 'q'), ('q', 'p', 'r'),
('r', 'q', 's'), ('s', 'p', 't'), ('t', 'q', 'u'),
('u', 'p', 'v'), ('v', 'q', 'w'), ('x', 'r', 'y'),
('x', 's', 'y')]
entity_lst = sorted({e
for (e, _, _) in triples}
| {e
for (_, _, e) in triples})
predicate_lst = sorted({p for (_, p, _) in triples})
nb_entities = len(entity_lst)
nb_predicates = len(predicate_lst)
entity_to_index = {e: i for i, e in enumerate(entity_lst)}
predicate_to_index = {p: i for i, p in enumerate(predicate_lst)}
for st in ['min', 'concat']:
with torch.no_grad():
kernel = GaussianKernel()
entity_embeddings = nn.Embedding(nb_entities,
embedding_size * 2,
sparse=True)
predicate_embeddings = nn.Embedding(nb_predicates,
embedding_size * 2,
sparse=True)
fact_rel = torch.from_numpy(
np.array([predicate_to_index[p] for (_, p, _) in triples]))
fact_arg1 = torch.from_numpy(
np.array([entity_to_index[s] for (s, _, _) in triples]))
fact_arg2 = torch.from_numpy(
np.array([entity_to_index[o] for (_, _, o) in triples]))
facts = [fact_rel, fact_arg1, fact_arg2]
model = NeuralKB(entity_embeddings=entity_embeddings,
predicate_embeddings=predicate_embeddings,
kernel=kernel,
facts=facts,
scoring_type=st)
indices = torch.from_numpy(
np.array([predicate_to_index['p'], predicate_to_index['q']]))
reformulator = SymbolicReformulator(predicate_embeddings, indices)
k = 5
rhoppy0 = RecursiveHoppy(model,
entity_embeddings,
hops=reformulator,
depth=0,
k=k)
rhoppy1 = RecursiveHoppy(model,
entity_embeddings,
hops=reformulator,
depth=1,
k=k)
rhoppy2 = RecursiveHoppy(model,
entity_embeddings,
hops=reformulator,
depth=2,
k=k)
rhoppy3 = RecursiveHoppy(model,
entity_embeddings,
hops=reformulator,
depth=3,
k=k)
rhoppy4 = RecursiveHoppy(model,
entity_embeddings,
hops=reformulator,
depth=4,
k=k)
xs_np = rs.randint(nb_entities, size=12)
xp_np = rs.randint(nb_predicates, size=12)
xo_np = rs.randint(nb_entities, size=12)
xs_np[0] = entity_to_index['a']
xp_np[0] = predicate_to_index['r']
xo_np[0] = entity_to_index['c']
xs_np[1] = entity_to_index['a']
xp_np[1] = predicate_to_index['r']
xo_np[1] = entity_to_index['e']
xs_np[2] = entity_to_index['a']
xp_np[2] = predicate_to_index['r']
xo_np[2] = entity_to_index['g']
xs_np[3] = entity_to_index['a']
xp_np[3] = predicate_to_index['r']
xo_np[3] = entity_to_index['i']
xs_np[4] = entity_to_index['a']
xp_np[4] = predicate_to_index['r']
xo_np[4] = entity_to_index['m']
xs_np[5] = entity_to_index['a']
xp_np[5] = predicate_to_index['r']
xo_np[5] = entity_to_index['o']
xs_np[6] = entity_to_index['a']
xp_np[6] = predicate_to_index['r']
xo_np[6] = entity_to_index['q']
xs_np[7] = entity_to_index['a']
xp_np[7] = predicate_to_index['r']
xo_np[7] = entity_to_index['s']
xs_np[8] = entity_to_index['a']
xp_np[8] = predicate_to_index['r']
xo_np[8] = entity_to_index['u']
# xs_np[9] = entity_to_index['a']
# xp_np[9] = predicate_to_index['r']
# xo_np[9] = entity_to_index['w']
xs = torch.from_numpy(xs_np)
xp = torch.from_numpy(xp_np)
xo = torch.from_numpy(xo_np)
xs_emb = entity_embeddings(xs)
xp_emb = predicate_embeddings(xp)
xo_emb = entity_embeddings(xo)
scores0 = rhoppy0.forward(xp_emb, xs_emb, xo_emb)
inf0 = rhoppy0.score(xp_emb, xs_emb, xo_emb)
for i in range(xs.shape[0]):
scores_sp, scores_po = scores0
inf_np = inf0.cpu().numpy()
scores_sp_np = scores_sp.cpu().numpy()
scores_po_np = scores_po.cpu().numpy()
np.testing.assert_allclose(inf_np[i],
scores_sp_np[i, xo[i]],
rtol=1e-5,
atol=1e-5)
np.testing.assert_allclose(inf_np[i],
scores_po_np[i, xs[i]],
rtol=1e-5,
atol=1e-5)
scores1 = rhoppy1.forward(xp_emb, xs_emb, xo_emb)
inf1 = rhoppy1.score(xp_emb, xs_emb, xo_emb)
for i in range(xs.shape[0]):
scores_sp, scores_po = scores1
inf_np = inf1.cpu().numpy()
scores_sp_np = scores_sp.cpu().numpy()
scores_po_np = scores_po.cpu().numpy()
np.testing.assert_allclose(inf_np[i],
scores_sp_np[i, xo[i]],
rtol=1e-5,
atol=1e-5)
np.testing.assert_allclose(inf_np[i],
scores_po_np[i, xs[i]],
rtol=1e-5,
atol=1e-5)
scores2 = rhoppy2.forward(xp_emb, xs_emb, xo_emb)
inf2 = rhoppy2.score(xp_emb, xs_emb, xo_emb)
for i in range(xs.shape[0]):
scores_sp, scores_po = scores2
inf_np = inf2.cpu().numpy()
scores_sp_np = scores_sp.cpu().numpy()
scores_po_np = scores_po.cpu().numpy()
np.testing.assert_allclose(inf_np[i],
scores_sp_np[i, xo[i]],
rtol=1e-1,
atol=1e-1)
np.testing.assert_allclose(inf_np[i],
scores_po_np[i, xs[i]],
rtol=1e-1,
atol=1e-1)
scores3 = rhoppy3.forward(xp_emb, xs_emb, xo_emb)
inf3 = rhoppy3.score(xp_emb, xs_emb, xo_emb)
for i in range(xs.shape[0]):
scores_sp, scores_po = scores3
inf_np = inf3.cpu().numpy()
scores_sp_np = scores_sp.cpu().numpy()
scores_po_np = scores_po.cpu().numpy()
np.testing.assert_allclose(inf_np[i],
scores_sp_np[i, xo[i]],
rtol=1e-1,
atol=1e-1)
np.testing.assert_allclose(inf_np[i],
scores_po_np[i, xs[i]],
rtol=1e-1,
atol=1e-1)
scores4 = rhoppy4.forward(xp_emb, xs_emb, xo_emb)
inf4 = rhoppy4.score(xp_emb, xs_emb, xo_emb)
for i in range(xs.shape[0]):
scores_sp, scores_po = scores4
inf_np = inf4.cpu().numpy()
scores_sp_np = scores_sp.cpu().numpy()
scores_po_np = scores_po.cpu().numpy()
np.testing.assert_allclose(inf_np[i],
scores_sp_np[i, xo[i]],
rtol=1e-1,
atol=1e-1)
np.testing.assert_allclose(inf_np[i],
scores_po_np[i, xs[i]],
rtol=1e-1,
atol=1e-1)
print(inf0)
print(inf1)
print(inf2)
print(inf3)
print(inf4)
inf0_np = inf0.cpu().numpy()
inf1_np = inf1.cpu().numpy()
inf2_np = inf2.cpu().numpy()
inf3_np = inf3.cpu().numpy()
inf4_np = inf4.cpu().numpy()
np.testing.assert_allclose(inf0_np,
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
rtol=1e-1,
atol=1e-1)
np.testing.assert_allclose(inf1_np,
[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
rtol=1e-1,
atol=1e-1)
np.testing.assert_allclose(inf2_np,
[1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
rtol=1e-1,
atol=1e-1)
np.testing.assert_allclose(inf3_np,
[1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
rtol=1e-1,
atol=1e-1)
np.testing.assert_allclose(inf4_np,
[1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0],
rtol=1e-1,
atol=1e-1)
def test_reasoning_v5():
torch.set_num_threads(multiprocessing.cpu_count())
nb_entities = 10
nb_predicates = 5
embedding_size = 10
rs = np.random.RandomState(0)
triples = [('a', 'p', 'b'), ('b', 'q', 'c'), ('c', 'r', 'd'),
('d', 's', 'e')]
entity_to_index = {'a': 0, 'b': 1, 'c': 2, 'd': 3, 'e': 4}
predicate_to_index = {'p': 0, 'q': 1, 'r': 2, 's': 3}
for st in ['min', 'concat']:
with torch.no_grad():
kernel = GaussianKernel()
entity_embeddings = nn.Embedding(nb_entities,
embedding_size * 2,
sparse=True)
predicate_embeddings = nn.Embedding(nb_predicates,
embedding_size * 2,
sparse=True)
fact_rel = torch.from_numpy(
np.array([predicate_to_index[p] for (_, p, _) in triples]))
fact_arg1 = torch.from_numpy(
np.array([entity_to_index[s] for (s, _, _) in triples]))
fact_arg2 = torch.from_numpy(
np.array([entity_to_index[o] for (_, _, o) in triples]))
facts = [fact_rel, fact_arg1, fact_arg2]
model = NeuralKB(entity_embeddings=entity_embeddings,
predicate_embeddings=predicate_embeddings,
kernel=kernel,
facts=facts,
scoring_type=st)
indices = torch.from_numpy(
np.array([
predicate_to_index['p'], predicate_to_index['q'],
predicate_to_index['r'], predicate_to_index['s']
]))
reformulator = SymbolicReformulator(predicate_embeddings, indices)
hoppy = SimpleHoppy(model, entity_embeddings, hops=reformulator)
rhoppy = RecursiveHoppy(model,
entity_embeddings,
hops=reformulator,
depth=1)
xs_np = rs.randint(nb_entities, size=32)
xp_np = rs.randint(nb_predicates, size=32)
xo_np = rs.randint(nb_entities, size=32)
xs_np[0] = 0
xp_np[0] = 0
xo_np[0] = 1
xs_np[1] = 1
xp_np[1] = 1
xo_np[1] = 2
xs_np[2] = 0
xp_np[2] = 3
xo_np[2] = 4
xs = torch.from_numpy(xs_np)
xp = torch.from_numpy(xp_np)
xo = torch.from_numpy(xo_np)
xs_emb = entity_embeddings(xs)
xp_emb = predicate_embeddings(xp)
xo_emb = entity_embeddings(xo)
scores = hoppy.forward(xp_emb, xs_emb, xo_emb)
inf = hoppy.score(xp_emb, xs_emb, xo_emb)
scores_h = rhoppy.depth_r_forward(xp_emb, xs_emb, xo_emb, depth=1)
inf_h = rhoppy.depth_r_score(xp_emb, xs_emb, xo_emb, depth=1)
print(inf)
print(inf_h)
assert inf[2] > 0.95
scores_sp, scores_po = scores
scores_h_sp, scores_h_po = scores_h
inf = inf.cpu().numpy()
scores_sp = scores_sp.cpu().numpy()
scores_po = scores_po.cpu().numpy()
inf_h = inf_h.cpu().numpy()
scores_h_sp = scores_h_sp.cpu().numpy()
scores_h_po = scores_h_po.cpu().numpy()
np.testing.assert_allclose(inf, inf_h)
np.testing.assert_allclose(scores_sp, scores_h_sp)
np.testing.assert_allclose(scores_po, scores_h_po)
for i in range(xs.shape[0]):
np.testing.assert_allclose(inf[i],
scores_sp[i, xo[i]],
rtol=1e-5,
atol=1e-5)
np.testing.assert_allclose(inf[i],
scores_po[i, xs[i]],
rtol=1e-5,
atol=1e-5)
np.testing.assert_allclose(inf_h[i],
scores_h_sp[i, xo[i]],
rtol=1e-5,
atol=1e-5)
np.testing.assert_allclose(inf_h[i],
scores_h_po[i, xs[i]],
rtol=1e-5,
atol=1e-5)