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 __init__(self, dimension):
super(Image_Transformer, self).__init__()
self.fc1 = MobiusLinear(2048, 1024)
self.mobius_relu = pm.mobius_fn_apply(F.relu, x)
self.fc2 = MobiusLinear(1024, dimension)