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