def mobius_gru_cell(
input: torch.Tensor,
hx: torch.Tensor,
weight_ih: torch.Tensor,
weight_hh: torch.Tensor,
bias: torch.Tensor,
c: torch.Tensor,
nonlin=None,
):
W_ir, W_ih, W_iz = weight_ih.chunk(3)
b_r, b_h, b_z = bias
W_hr, W_hh, W_hz = weight_hh.chunk(3)
z_t = pmath.logmap0(one_rnn_transform(W_hz, hx, W_iz, input, b_z, c),
c=c).sigmoid()
r_t = pmath.logmap0(one_rnn_transform(W_hr, hx, W_ir, input, b_r, c),
c=c).sigmoid()
rh_t = pmath.mobius_pointwise_mul(r_t, hx, c=c)
h_tilde = one_rnn_transform(W_hh, rh_t, W_ih, input, b_h, c)
if nonlin is not None:
h_tilde = pmath.mobius_fn_apply(nonlin, h_tilde, c=c)
delta_h = pmath.mobius_add(-hx, h_tilde, c=c)
h_out = pmath.mobius_add(hx,
pmath.mobius_pointwise_mul(z_t, delta_h, c=c),
c=c)
return h_out
def mobius_linear(
input,
weight,
bias=None,
hyperbolic_input=True,
hyperbolic_bias=True,
nonlin=None,
c=1.0,
):
if hyperbolic_input:
output = pmath.mobius_matvec(weight, input, c=c)
print('x')
print(weight.grad)
print('x')
else:
output = torch.nn.functional.linear(input, weight)
# output = rm_linear(input, weight)
output = pmath.expmap0(output, c=c)
if bias is not None:
if not hyperbolic_bias:
bias = pmath.expmap0(bias, c=c)
output = pmath.mobius_add(output, bias, c=c)
if nonlin is not None:
output = pmath.mobius_fn_apply(nonlin, output, c=c)
output = pmath.project(output, c=c)
return output
def one_rnn_transform(W, h, U, x, b, c):
W_otimes_h = pmath.mobius_matvec(W, h, c=c)
U_otimes_x = pmath.mobius_matvec(U, x, c=c)
Wh_plus_Ux = pmath.mobius_add(W_otimes_h, U_otimes_x, c=c)
return pmath.mobius_add(Wh_plus_Ux, b, c=c)
def forward(self, input):
source_input = input[0][0]
# print(source_input)
target_input = input[0][1]
alignment = input[1]
batch_size = alignment.shape[0]
source_input_data = self.embedding(source_input.data)
target_input_data = self.embedding(target_input.data)
zero_hidden = torch.zeros(self.num_layers,
batch_size,
self.hidden_dim,
device=self.device or source_input.device,
dtype=source_input_data.dtype)
# print(self.cell_type)
if self.embedding_type == "eucl" and "hyp" in self.cell_type: # This is for the example
# print(source_input_data.shape)
source_input_data = pmath.expmap0(source_input_data, c=self.c)
# print(source_input_data.shape)
target_input_data = pmath.expmap0(target_input_data, c=self.c)
elif self.embedding_type == "hyp" and "eucl" in self.cell_type:
source_input_data = pmath.logmap0(source_input_data, c=self.c)
target_input_data = pmath.logmap0(target_input_data, c=self.c)
# ht: (num_layers * num_directions, batch, hidden_size)
# print(source_input.batch_sizes.shape)
source_input = torch.nn.utils.rnn.PackedSequence(
source_input_data, source_input.batch_sizes)
target_input = torch.nn.utils.rnn.PackedSequence(
target_input_data, target_input.batch_sizes)
_, source_hidden = self.cell_source(source_input, zero_hidden)
_, target_hidden = self.cell_target(target_input, zero_hidden)
# take hiddens from the last layer
source_hidden = source_hidden[-1]
# print(target_hidden)
target_hidden = target_hidden[-1][alignment]
# print(alignment)
# print(target_hidden)
if self.decision_type == "hyp":
if "eucl" in self.cell_type:
source_hidden = pmath.expmap0(source_hidden, c=self.c)
target_hidden = pmath.expmap0(target_hidden, c=self.c)
source_projected = self.projector_source(source_hidden)
target_projected = self.projector_target(target_hidden)
projected = pmath.mobius_add(source_projected,
target_projected,
c=self.ball.c)
if self.use_distance_as_feature:
dist = (pmath.dist(source_hidden,
target_hidden,
dim=-1,
keepdim=True,
c=self.ball.c)**2)
bias = pmath.mobius_scalar_mul(dist,
self.dist_bias,
c=self.ball.c)
projected = pmath.mobius_add(projected, bias, c=self.ball.c)
else:
if "hyp" in self.cell_type:
source_hidden = pmath.logmap0(source_hidden, c=self.c)
target_hidden = pmath.logmap0(target_hidden, c=self.c)
projected = self.projector(
torch.cat((source_hidden, target_hidden), dim=-1))
if self.use_distance_as_feature:
dist = torch.sum((source_hidden - target_hidden).pow(2),
dim=-1,
keepdim=True)
bias = self.dist_bias * dist
projected = projected + bias
logits = self.logits(projected)
# CrossEntropy accepts logits
return logits