def up_reduce_func(self, nodes):
bs = nodes.mailbox['beta_ch'].shape[0]
n_ch = nodes.mailbox['beta_ch'].shape[1]
beta_ch = thlp.zeros(bs, self.max_output_degree, self.h_size)
beta_ch[:, :n_ch, :] = nodes.mailbox['beta_ch']
# compute beta_r
U = self.__gather_param__(
self.U,
types=nodes.data['t'] if self.num_types > 1 else None,
)
# U has shape bs x L x h x rank
gamma_ch_rl = thlp.sum_over(thlp.mul(U, beta_ch.unsqueeze(3)),
2) # has shape bs x L x rank
G = self.__gather_param__(
self.G, types=nodes.data['t'] if self.num_types > 1 else None)
for i in range(self.max_output_degree):
# TODO: we are assuming last are bottom
if i < n_ch:
btm = thlp.zeros(bs, 1)
x = th.cat((gamma_ch_rl[:, i, :], btm), 1)
else:
x = thlp.zeros(1, self.rank + 1)
x[:, -1] = 0
new_shape = [x.shape[0]] + [1] * i + [
self.rank + 1
] + [1] * (self.max_output_degree - i)
G = thlp.mul(G, x.view(*new_shape))
gamma_r = thlp.sum_over(G, list(range(1, self.max_output_degree + 1)))
return {'gamma_r': gamma_r, 'gamma_ch_all': G, 'beta_ch': beta_ch}
def up_reduce_func(self, nodes):
bs = nodes.mailbox['beta_ch'].shape[0]
n_ch = nodes.mailbox['beta_ch'].shape[1]
beta_ch = thlp.zeros(bs, self.max_output_degree, self.h_size)
beta_ch[:, :n_ch, :] = nodes.mailbox['beta_ch']
# TODO: we are assuming beta_ch is ordered accoridng pos. It allows bottom at the end
# compute beta_r
U = self.__gather_param__(
self.U, types=nodes.data['t'] if self.num_types > 1 else None)
# U has shape bs x L x h x rank+1 x rank
gamma_ch_rl = thlp.sum_over(
thlp.mul(U,
beta_ch.unsqueeze(3).unsqueeze(4)), 2)
# has shape bs x L x rank+1 x rank
gamma_less_l = thlp.zeros(bs, self.max_output_degree, self.rank + 1)
gamma_less_l[:, 0, :-1] = gamma_ch_rl[:, 0,
-1, :] # has shape bs x rank
for i in range(1, n_ch):
gamma_i = gamma_ch_rl[:, i, :, :] # has shape bs x (rank+1) x rank
gamma_prev = gamma_less_l[:, i - 1, :].unsqueeze(
2) # has shape bs x (rank+1) x 1
gamma_less_l[:,
i, :-1] = thlp.sum_over(thlp.mul(gamma_i, gamma_prev),
1)
gamma_less_l[:, n_ch:, -1] = 0
# R_out = self.__gather_param__(self.R_output, types=nodes.data['t'] if self.num_types > 1 else None)
# gamma_r = thlp.sum_over(thlp.mul(R_out, gamma_less_l[:, n_ch-1, :-1].unsqueeze(2)), 1) # has shape bs x rank
gamma_r = gamma_less_l[:, n_ch - 1, :-1]
return {
'gamma_less_l': gamma_less_l,
'gamma_r': gamma_r,
'beta_ch': beta_ch,
'n_ch': th.full([bs], n_ch, dtype=th.long)
}
def up_reduce_func(self, nodes):
beta_ch = nodes.mailbox['beta_ch'] # has shape (bs x n_ch x h)
n_ch = beta_ch.shape[1]
bs = beta_ch.shape[0]
U = self.__gather_param__(
self.U, types=nodes.data['t'] if self.num_types > 1 else None)
for i in range(self.max_output_degree):
# TODO: we are assuming last are bottom
if i < n_ch:
btm = thlp.zeros(bs, 1)
x = th.cat((beta_ch[:, i, :], btm), 1)
else:
x = thlp.zeros(1, self.h_size + 1)
x[:, -1] = 0
new_shape = [x.shape[0]] + [1] * i + [
self.h_size + 1
] + [1] * (self.max_output_degree - i)
U = thlp.mul(U, x.view(*new_shape))
beta = thlp.sum_over(U, list(range(1, self.max_output_degree + 1)))
return {'beta_np': beta, 'beta_ch': U}
def up_reduce_func(self, nodes):
bs = nodes.mailbox['beta_ch'].shape[0]
n_ch = nodes.mailbox['beta_ch'].shape[1]
beta_ch = thlp.zeros(bs, self.max_output_degree, self.h_size)
beta_ch[:, :n_ch, :] = nodes.mailbox['beta_ch']
U = self.__gather_param__(
self.U, types=nodes.data['t'] if self.num_types > 1 else None)
# has shape bs x L x h x rank
gamma_ch = thlp.mul(beta_ch.unsqueeze(3),
U) # has shape (bs x L x h x rank)
gamma_p_ch = thlp.sum_over(gamma_ch, 2) # has shape (bs x L x rank)
gamma_r = th.sum(gamma_p_ch[:, :n_ch, :], 1) # has shape (bs x rank)
return {
'gamma_r': gamma_r,
'gamma_ch': gamma_ch,
'gamma_p_ch': gamma_p_ch
}
def up_apply_node_func(self, nodes):
x = nodes.data['evid'] # represents P(x_u | Q_u) have size bs x h
if 'gamma_r' in nodes.data:
U_out = self.__gather_param__(
self.U_output,
types=nodes.data['t'] if self.num_types > 1 else None)
# U_out has shape bs x rank x h
beta = thlp.sum_over(
thlp.mul(nodes.data['gamma_r'].unsqueeze(2), U_out), 1)
gamma_r = nodes.data['gamma_r']
gamma_ch_all = nodes.data['gamma_ch_all']
beta_ch = nodes.data['beta_ch']
else:
beta = self.__gather_param__(
self.p,
types=nodes.data['t'] if self.num_types > 1 else None,
pos=nodes.data['pos'])
bs = beta.size(0)
gamma_r = th.zeros((bs, self.rank))
gamma_ch_all = th.zeros(
[bs] + [self.rank + 1
for i in range(self.max_output_degree)] + [self.rank])
beta_ch = thlp.zeros(bs, self.max_output_degree, self.h_size)
beta = thlp.mul(x, beta) # has shape (bs x h)
# normalise
beta, N_u = thlp.normalise(beta, 1, get_Z=True)
return {
'beta': beta,
'N_u': N_u,
'gamma_r': gamma_r,
'gamma_ch_all': gamma_ch_all,
'beta_ch': beta_ch
}
def down_apply_node_func(self, nodes):
if 'eta' in nodes.data:
eta_u = nodes.data['eta']
else:
# root
eta_u = nodes.data['beta']
is_leaf = nodes.data['is_leaf']
is_internal = th.logical_not(is_leaf)
n_ch_list = nodes.data['n_ch'][is_internal] - 1
gamma_r = nodes.data['gamma_r'][is_internal] # has shape (bs x rank)
# has shape bs x rank+1
gamma_less_l = nodes.data['gamma_less_l'][
is_internal] # has shape bs x L x rank+1
beta_ch = nodes.data['beta_ch'][is_internal] # has shape bs x L x h
t = nodes.data['t'][is_internal]
self.accumulate_posterior(
self.p,
eta_u[is_leaf],
types=nodes.data['t'][is_leaf] if self.num_types > 1 else None,
pos=nodes.data['pos'][is_leaf]
if not self.pos_stationarity else None)
eta_u_ch_all = thlp.zeros(eta_u.shape[0], self.max_output_degree,
self.h_size)
if th.any(is_internal):
# computation only on internal nodes
# compute P(Q_u, R_u | X)
U_out = self.__gather_param__(
self.U_output, types=t if self.num_types > 1 else None)
# U_out has shape bs x rank x h
a = thlp.mul(gamma_r.unsqueeze(2), U_out)
b = thlp.sum_over(a, 1, keepdim=True)
eta_ur = thlp.div(thlp.mul(a, eta_u[is_internal].unsqueeze(1)),
b) # has shape bs x rank x h
# compute P(R_u, R_L | X)
eta_r = thlp.sum_over(eta_ur, 2) # has shape bs x rank_U
# R_out = self.__gather_param__(self.R_output, types=t if self.num_types > 1 else None)
# # R_out has shape bs x rank_L x rank_U
# a = thlp.mul(R_out, gamma_L[:, :-1].unsqueeze(2)) # has shape bs x rank_L x rank_U
# b = thlp.sum_over(a, 1, keepdim=True)
# eta_rul = thlp.div(thlp.mul(a, eta_r.unsqueeze(1)), b) # has shape bs x rank_L x rank_U
# compute P(R_l, R_l-1, Q_l | X)
# eta_rL = thlp.sum_over(eta_rul, 2) # has shape bs x rank
eta_rL = eta_r
U = self.__gather_param__(self.U,
types=t if self.num_types > 1 else None)
if self.pos_stationarity:
U = U.expand((-1, self.max_output_degree, -1, -1, -1))
if U.size(0) == 1:
U = U.expand((gamma_less_l.size(0), -1, -1, -1, -1))
# U has shape bs x L x h x rank+1 x rank
eta_u_ch = thlp.zeros(eta_rL.shape[0], self.max_output_degree,
self.h_size)
last_eta = eta_rL # has shape bs x rank
for i in range(self.max_output_degree - 1, -1, -1):
pos_flag = i <= n_ch_list
if th.any(pos_flag):
if i > 0:
a = thlp.mul(
U[pos_flag, i, :, :, :],
gamma_less_l[pos_flag,
i - 1, :].unsqueeze(1).unsqueeze(3))
a = thlp.mul(
a, beta_ch[pos_flag,
i, :].unsqueeze(2).unsqueeze(3))
else:
a = thlp.zeros(*(U.shape[:1] + U.shape[2:]))
a[:, :,
-1, :] = thlp.mul(U[:, i, :, -1, :],
beta_ch[:, i, :].unsqueeze(2))
b = thlp.sum_over(a, (1, 2), keepdim=True)
eta_rul_rlprec = thlp.div(
thlp.mul(
a,
last_eta[pos_flag, :].unsqueeze(1).unsqueeze(2)),
b)
self.accumulate_posterior(
self.U,
eta_rul_rlprec,
types=t[pos_flag] if self.num_types > 1 else None,
pos=th.full(
(eta_rul_rlprec.shape[0],
1), i, dtype=th.long).squeeze(1)
if not self.pos_stationarity else None)
eta_u_ch[pos_flag,
i, :] = thlp.sum_over(eta_rul_rlprec, (2, 3))
last_eta[pos_flag, :] = thlp.sum_over(
eta_rul_rlprec, (1, 3))[:, :-1]
# accumulate posterior
self.accumulate_posterior(self.U_output,
eta_ur,
types=t if self.num_types > 1 else None)
# self.accumulate_posterior(self.R_output, eta_rul, types=t if self.num_types > 1 else None)
eta_u_ch_all[is_internal] = eta_u_ch
return {'eta_ch': eta_u_ch_all, 'eta': eta_u}
def down_apply_node_func(self, nodes):
if 'eta' in nodes.data:
eta_u = nodes.data['eta']
else:
# root
eta_u = nodes.data['beta']
bs = eta_u.shape[0]
gamma_r = nodes.data['gamma_r'] # has shape (bs x rank)
U_out = self.__gather_param__(
self.U_output,
types=nodes.data['t'] if self.num_types > 1 else None)
# U_out has shape bs x rank x h
a = thlp.mul(gamma_r.unsqueeze(2), U_out)
b = thlp.sum_over(a, 1, keepdim=True)
# P(Q_u, R_u | X)
eta_ur = thlp.div(thlp.mul(a, eta_u.unsqueeze(1)),
b) # has shape bs x rank x h
eta_r = thlp.sum_over(eta_ur, 2) # has shape bs x rank
gamma_ch_all = nodes.data[
'gamma_ch_all'] # has shape bs x r x ... x r x r
new_shape = [-1] + [1] * self.max_output_degree + [self.rank]
# P(R_u, R_1, ..., R_L | X)
eta_ru_rch = thlp.div(thlp.mul(gamma_ch_all, eta_r.view(*new_shape)),
gamma_r.view(*new_shape))
# P(R_1, ..., R_L | X)
eta_rch = thlp.sum_over(eta_ru_rch, -1) # has shape bs x r+1 x ... r+1
eta_rl = thlp.zeros(bs, self.max_output_degree, self.rank + 1)
for i in range(self.max_output_degree):
sum_over_var = list(
set(range(1, self.max_output_degree + 1)) - {i + 1})
eta_rl[:, i, :] = thlp.sum_over(eta_rch, sum_over_var)
U = self.__gather_param__(
self.U, types=nodes.data['t'] if self.num_types > 1 else None)
# U has shape bs x L x h x rank
a = thlp.mul(U, nodes.data['beta_ch'].unsqueeze(3))
b = thlp.sum_over(a, 2, keepdim=True)
# P(Q_l, R_l | X)
eta_rql = thlp.div(thlp.mul(a, eta_rl[:, :, :-1].unsqueeze(2)),
b) # has shape bs x L x h x rank
# accumulate posterior
is_leaf = nodes.data['is_leaf']
is_internal = th.logical_not(is_leaf)
self.accumulate_posterior(
self.U,
eta_rql[is_internal],
types=nodes.data['t'][is_internal] if self.num_types > 1 else None)
self.accumulate_posterior(
self.U_output,
eta_ur[is_internal],
types=nodes.data['t'][is_internal] if self.num_types > 1 else None)
self.accumulate_posterior(
self.G,
eta_ru_rch[is_internal],
types=nodes.data['t'][is_internal] if self.num_types > 1 else None)
self.accumulate_posterior(
self.p,
eta_u[is_leaf],
types=nodes.data['t'][is_leaf] if self.num_types > 1 else None,
pos=nodes.data['pos'][is_leaf])
return {'eta_ch': thlp.sum_over(eta_rql, 3), 'eta': eta_u}