def mobius_linear(
input,
weight,
bias=None,
hyperbolic_input=True,
hyperbolic_bias=True,
nonlin=None,
k=-1.0,
):
k = torch.tensor(k)
if hyperbolic_input:
output = mobius_matvec(weight, input, k=k)
else:
output = torch.nn.functional.linear(input, weight)
output = gmath.expmap0(output, k=k)
if bias is not None:
if not hyperbolic_bias:
bias = gmath.expmap0(bias, k=k)
output = gmath.mobius_add(output,
bias.unsqueeze(0).expand_as(output),
k=k)
if nonlin is not None:
output = gmath.mobius_fn_apply(nonlin, output, k=k)
output = gmath.project(output, k=k)
return output
def mobius_gru_cell(
input: torch.Tensor,
hx: torch.Tensor,
weight_ih: torch.Tensor,
weight_hh: torch.Tensor,
bias: torch.Tensor,
k: 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 = gmath.logmap0(one_rnn_transform(W_hz, hx, W_iz, input, b_z, k),
k=k).sigmoid()
r_t = gmath.logmap0(one_rnn_transform(W_hr, hx, W_ir, input, b_r, k),
k=k).sigmoid()
rh_t = gmath.mobius_pointwise_mul(r_t, hx, k=k)
h_tilde = one_rnn_transform(W_hh, rh_t, W_ih, input, b_h, k)
if nonlin is not None:
h_tilde = gmath.mobius_fn_apply(nonlin, h_tilde, k=k)
delta_h = gmath.mobius_add(-hx, h_tilde, k=k)
h_out = gmath.mobius_add(hx,
gmath.mobius_pointwise_mul(z_t, delta_h, k=k),
k=k)
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, k=to_k(c, weight))
else:
output = torch.nn.functional.linear(input, weight)
output = pmath.expmap0(output, k=to_k(c, output))
if bias is not None:
if not hyperbolic_bias:
bias = pmath.expmap0(bias, k=to_k(c, bias))
output = pmath.mobius_add(output, bias, k=to_k(c, output))
if nonlin is not None:
output = pmath.mobius_fn_apply(nonlin, output, k=to_k(c, output))
output = pmath.project(output, k=to_k(c, output))
return output