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 one_rnn_transform(W, h, U, x, b, c): W_otimes_h = pmath.mobius_matvec(W, h, k=to_k(c, W)) U_otimes_x = pmath.mobius_matvec(U, x, k=to_k(c, U)) Wh_plus_Ux = pmath.mobius_add(W_otimes_h, U_otimes_x, k=to_k(c, W_otimes_h)) return pmath.mobius_add(Wh_plus_Ux, b, k=to_k(c, Wh_plus_Ux))
def transition(self, x, h): """ :param x: batch x input :param h: hidden x hidden :return: batch x hidden """ W_otimes_h = pmath.mobius_matvec(self.w, h, k=self.ball.k) U_otimes_x = pmath.mobius_matvec(self.u, x, k=self.ball.k) Wh_plus_Ux = pmath.mobius_add(W_otimes_h, U_otimes_x, k=self.ball.k) return pmath.mobius_add(Wh_plus_Ux, self.b, k=self.ball.k)
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 exp_map_x_polar(x: Tensor, radius: Tensor, v: Tensor, c: Tensor) -> Tensor: v = v + eps # Perturb v to avoid dealing with v = 0 norm_v = torch.norm(v, p=2, dim=-1, keepdim=True) assert len(radius.shape) == len(v.shape) assert radius.shape[:-1] == v.shape[:-1] second_term = (torch.tanh(c.sqrt() * radius / 2) / (c.sqrt() * norm_v)) * v assert x.shape == second_term.shape return pm.mobius_add(x, second_term, c=c)
def poincare_distance_c(x: Tensor, y: Tensor, c: Tensor, keepdim: bool = True, **kwargs: Any) -> Tensor: # res = pm.dist(x, y, c=c, keepdim=keepdim, **kwargs) sqrt_c = sqrt(c) mob = pm.mobius_add(-x, y, c=c, dim=-1).norm(dim=-1, p=2, keepdim=keepdim) arg = sqrt_c * mob dist_c = atanh(arg) res = dist_c * 2 / sqrt_c assert torch.isfinite(res).all() return res
def _dist2plane(self, x, a, p, c, k, keepdim: bool = False, signed: bool = False, dim: int = -1): """ Taken from geoopt and corrected so it returns a_norm and this value does not have to be calculated twice """ sqrt_c = c ** 0.5 minus_p_plus_x = pmath.mobius_add(-p, x, k=k, dim=dim) mpx_sqnorm = minus_p_plus_x.pow(2).sum(dim=dim, keepdim=keepdim).clamp_min(MIN_NORM) mpx_dot_a = (minus_p_plus_x * a).sum(dim=dim, keepdim=keepdim) if not signed: mpx_dot_a = mpx_dot_a.abs() a_norm = a.norm(dim=dim, keepdim=keepdim, p=2).clamp_min(MIN_NORM) num = 2 * sqrt_c * mpx_dot_a denom = (1 - c * mpx_sqnorm) * a_norm return pmath.arsinh(num / denom.clamp_min(MIN_NORM)) / sqrt_c, a_norm
def __init__(self, args, pos_embeds_rows, input_dims): super(DistanceAttention, self).__init__() # pos embeds self.manifold = gt.PoincareBall( ) if args.attn_metric == cs.HY else gt.Euclidean() with torch.no_grad(): pos_embeds = init_embeddings(pos_embeds_rows, input_dims) if args.attn_metric == cs.HY: pos_embeds = pmath.expmap0(pos_embeds, k=self.manifold.k) beta = torch.Tensor(1).uniform_(-0.01, 0.01) self.position_embeds = gt.ManifoldParameter(pos_embeds, manifold=self.manifold) if args.attn_metric == cs.HY: self.key_dense = hnn.MobiusLinear(input_dims, input_dims, hyperbolic_input=True, hyperbolic_bias=True) self.query_dense = hnn.MobiusLinear(input_dims, input_dims, hyperbolic_input=True, hyperbolic_bias=True) self.addition = lambda x, y: pmath.mobius_add( x, y, k=self.manifold.k) self.distance_function = lambda x, y: pmath.dist( x, y, k=self.manifold.k) self.midpoint = hnn.weighted_mobius_midpoint else: self.key_dense = nn.Linear(input_dims, input_dims) self.query_dense = nn.Linear(input_dims, input_dims) self.addition = torch.add self.distance_function = utils.euclidean_distance self.midpoint = hnn.weighted_euclidean_midpoint self.encoder_to_attn_map = define_mapping(args.encoder_metric, args.attn_metric, args.c) self.attention_function = nn.Softmax( dim=1) if args.attn == "softmax" else nn.Sigmoid() self.beta = torch.nn.Parameter(beta, requires_grad=True)
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
def one_rnn_transform(W, h, U, x, b, k): W_otimes_h = gmath.mobius_matvec(W, h, k=k) U_otimes_x = gmath.mobius_matvec(U, x, k=k) Wh_plus_Ux = gmath.mobius_add(W_otimes_h, U_otimes_x, k=k) return gmath.mobius_add(Wh_plus_Ux, b, k=k)
def forward(self, input): source_input = input[0] target_input = input[1] alignment = input[2] 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) if self.embedding_type == "eucl" and "hyp" in self.cell_type: source_input_data = pmath.expmap0(source_input_data, c=self.c) 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) 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] target_hidden = target_hidden[-1][alignment] 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
def add(self, a, b): out = pmath.mobius_add(a, b, k=self.ball.k) return pmath.project(out, k=self.ball.k)