def test_compare_manifolds():
m1 = geoopt.Euclidean()
m2 = geoopt.Euclidean(ndim=1)
tensor = geoopt.ManifoldTensor(10, manifold=m1)
with pytest.raises(ValueError) as e:
_ = geoopt.ManifoldParameter(tensor, manifold=m2)
assert e.match("Manifolds do not match")
def product_case():
torch.manual_seed(42)
ex = [
torch.randn(10),
torch.randn(3) / 10,
torch.randn(3, 2),
torch.randn(())
]
ev = [
torch.randn(10),
torch.randn(3) / 10,
torch.randn(3, 2),
torch.randn(())
]
manifolds = [
geoopt.Sphere(),
geoopt.PoincareBall(),
geoopt.Stiefel(),
geoopt.Euclidean(),
]
x = [manifolds[i].projx(ex[i]) for i in range(len(manifolds))]
v = [manifolds[i].proju(x[i], ev[i]) for i in range(len(manifolds))]
product_manifold = geoopt.ProductManifold(*((manifolds[i], ex[i].shape)
for i in range(len(ex))))
yield UnaryCase(
manifold_shapes[geoopt.ProductManifold],
product_manifold.pack_point(*x),
product_manifold.pack_point(*ex),
product_manifold.pack_point(*v),
product_manifold.pack_point(*ev),
product_manifold,
)
# + 1 case without stiefel
torch.manual_seed(42)
ex = [torch.randn(10), torch.randn(3) / 10, torch.randn(())]
ev = [torch.randn(10), torch.randn(3) / 10, torch.randn(())]
manifolds = [
geoopt.Sphere(),
geoopt.PoincareBall(),
# geoopt.Stiefel(),
geoopt.Euclidean(),
]
x = [manifolds[i].projx(ex[i]) for i in range(len(manifolds))]
v = [manifolds[i].proju(x[i], ev[i]) for i in range(len(manifolds))]
product_manifold = geoopt.ProductManifold(*((manifolds[i], ex[i].shape)
for i in range(len(ex))))
yield UnaryCase(
manifold_shapes[geoopt.ProductManifold],
product_manifold.pack_point(*x),
product_manifold.pack_point(*ex),
product_manifold.pack_point(*v),
product_manifold.pack_point(*ev),
product_manifold,
)
def __init__(self,
num_embeddings,
embedding_dim,
manifold=geoopt.Euclidean(),
_weight=None):
super(LookupEmbedding, self).__init__()
if isinstance(embedding_dim, int):
embedding_dim = (embedding_dim, )
self.num_embeddings = num_embeddings
self.embedding_dim = embedding_dim
self.manifold = manifold
if _weight is None:
_weight = torch.Tensor(num_embeddings, *embedding_dim)
self.weight = geoopt.ManifoldParameter(_weight,
manifold=self.manifold)
self.reset_parameters()
else:
assert _weight.shape == (
num_embeddings,
*embedding_dim,
), "_weight MUST be of shape (num_embeddings, *embedding_dim)"
self.weight = geoopt.ManifoldParameter(_weight,
manifold=self.manifold)
def test_ndim_expected_behaviour(ndim):
eucl = geoopt.Euclidean(ndim=ndim)
point = eucl.random_normal(2, 2, 2, 2, 2)
point1 = eucl.random_normal(2, 2, 2, 2, 2)
# since spaces are the same we just use same method in test
tangent = eucl.random_normal(2, 2, 2, 2, 2)
inner = eucl.inner(point, tangent, tangent)
assert inner.dim() == tangent.dim() - ndim
inner = eucl.inner(point, tangent)
assert inner.dim() == tangent.dim() - ndim
norm = eucl.norm(point, tangent)
assert norm.dim() == tangent.dim() - ndim
dist = eucl.dist(point, point1)
assert dist.dim() == point.dim() - ndim
# keepdim now
inner = eucl.inner(point, tangent, tangent, keepdim=True)
assert inner.dim() == tangent.dim()
inner = eucl.inner(point, tangent, keepdim=True)
assert inner.dim() == tangent.dim()
norm = eucl.norm(point, tangent, keepdim=True)
assert norm.dim() == tangent.dim()
dist = eucl.dist(point, point1, keepdim=True)
assert dist.dim() == point.dim()
class ManifoldFactory:
geoopt_manifolds = {
"euclidean": lambda dims: gt.Euclidean(1),
"poincare": lambda dims: gt.PoincareBall(),
"lorentz": lambda dims: gt.Lorentz(),
"sphere": lambda dims: gt.Sphere(),
"prod-hysph": get_prod_hysph_manifold,
"prod-hyhy": get_prod_hyhy_manifold,
"prod-hyeu": get_prod_hyeu_manifold,
"prod-sphsph": get_prod_sphsph_manifold,
"spd": lambda dims: SymmetricPositiveDefinite(),
}
sympa_manifolds = {
"upper": UpperHalfManifold,
"bounded": BoundedDomainManifold,
"dual": CompactDualManifold
}
@classmethod
def get_manifold(cls, manifold_name, metric_name, dims):
if manifold_name in cls.geoopt_manifolds:
return cls.geoopt_manifolds[manifold_name](dims)
manifold = cls.sympa_manifolds[manifold_name]
metric = MetricType.from_str(metric_name)
return manifold(dims=dims, metric=metric)
def __init__(self, input_size, hidden_size, num_layers=1, bias=True, nonlin=None):
super().__init__()
self.manifold = gt.Euclidean()
self.input_size = input_size
self.hidden_size = hidden_size
self.num_layers = num_layers
self.bias = bias
self.weight_ih = torch.nn.ParameterList(
[torch.nn.Parameter(torch.Tensor(3 * hidden_size, input_size if i == 0 else hidden_size))
for i in range(num_layers)]
)
self.weight_hh = torch.nn.ParameterList(
[torch.nn.Parameter(torch.Tensor(3 * hidden_size, hidden_size)) for _ in range(num_layers)]
)
if bias:
biases = []
for i in range(num_layers):
bias = torch.randn(3, hidden_size) * 1e-5
bias = gt.ManifoldParameter(bias, manifold=self.manifold)
biases.append(bias)
self.bias = torch.nn.ParameterList(biases)
else:
self.register_buffer("bias", None)
self.nonlin = nonlin
self.reset_parameters()
def test_product():
manifold = geoopt.ProductManifold(
(geoopt.Sphere(), 10),
(geoopt.PoincareBall(), 3),
(geoopt.Stiefel(), (20, 2)),
(geoopt.Euclidean(), 43),
)
sample = manifold.random(20, manifold.n_elements)
manifold.assert_check_point_on_manifold(sample)
def euclidean_case():
torch.manual_seed(42)
shape = manifold_shapes[geoopt.manifolds.Euclidean]
ex = torch.randn(*shape, dtype=torch.float64)
ev = torch.randn(*shape, dtype=torch.float64)
x = ex.clone()
v = ev.clone()
manifold = geoopt.Euclidean(ndim=1)
x = geoopt.ManifoldTensor(x, manifold=manifold)
case = UnaryCase(shape, x, ex, v, ev, manifold)
yield case
def __init__(self, output_dim, input_dims, second_input_dim=None, third_input_dim=None, nonlin=None):
super(EuclConcat, self).__init__()
b_input_dims = second_input_dim if second_input_dim is not None else input_dims
self.lin_a = nn.Linear(input_dims, output_dim, bias=False)
self.lin_b = nn.Linear(b_input_dims, output_dim, bias=False)
if third_input_dim:
self.lin_c = nn.Linear(third_input_dim, output_dim, bias=False)
self.manifold = gt.Euclidean()
self.bias = gt.ManifoldParameter(torch.randn(output_dim) * 1e-5, manifold=self.manifold)
self.nonlin = nonlin if nonlin is not None else lambda x: x
def __init__(self, input_size, hidden_size):
super(EuclRNN, self).__init__()
self.manifold = gt.Euclidean()
self.input_size = input_size
self.hidden_size = hidden_size
# k = (1 / hidden_size)**0.5
k_w = (6 / (self.hidden_size + self.hidden_size)) ** 0.5 # xavier uniform
k_u = (6 / (self.input_size + self.hidden_size)) ** 0.5 # xavier uniform
self.w = gt.ManifoldParameter(gt.ManifoldTensor(hidden_size, hidden_size).uniform_(-k_w, k_w))
self.u = gt.ManifoldParameter(gt.ManifoldTensor(input_size, hidden_size).uniform_(-k_u, k_u))
bias = torch.randn(hidden_size) * 1e-5
self.b = gt.ManifoldParameter(bias, manifold=self.manifold)
def test_ismanifold():
m1 = geoopt.Euclidean()
assert geoopt.ismanifold(m1, geoopt.Euclidean)
m1 = geoopt.Scaled(m1)
m1 = geoopt.Scaled(m1)
assert geoopt.ismanifold(m1, geoopt.Euclidean)
with pytest.raises(TypeError):
geoopt.ismanifold(m1, int)
with pytest.raises(TypeError):
geoopt.ismanifold(m1, 1)
assert not geoopt.ismanifold(1, geoopt.Euclidean)
def __init__(self, in_features, out_features):
"""
:param in_features: number of dimensions of the input
:param out_features: number of classes
"""
super().__init__()
self.in_features = in_features
self.out_features = out_features
self.manifold = gt.Euclidean()
points = torch.randn(out_features, in_features) * 1e-5
self.p_k = gt.ManifoldParameter(points, manifold=self.manifold)
tangent = torch.Tensor(out_features, in_features)
stdv = (6 / (out_features + in_features)) ** 0.5 # xavier uniform
torch.nn.init.uniform_(tangent, -stdv, stdv)
self.a_k = torch.nn.Parameter(tangent)
def __init__(self, args, vocabs, word2vec=None):
self.args = args
super(Model, self).__init__()
self.word_embed_manifold = gt.PoincareBall(
) if args.embedding_metric == cs.HY else gt.Euclidean()
self.train_word_embeds = args.train_word_embeds == 1
self.word_lut = self.init_lut(
word2vec, len(vocabs[cs.TOKEN_VOCAB].label2wordvec_idx),
args.word_emb_size)
self.word_lut.requires_grad = self.train_word_embeds
self.concat_dropout = nn.Dropout(p=args.concat_dropout)
self.classif_dropout = nn.Dropout(p=args.classif_dropout)
# encoders
self.mention_encoder = MentionEncoder(vocabs[cs.CHAR_VOCAB], args)
self.context_encoder = ContextEncoder(args)
men_dim = self.mention_encoder.mention_output_dim
char_dim = self.mention_encoder.char_output_dim
ctx_dim = self.context_encoder.ctx_output_dim
# ctx concat and attn
ctx_concat_layer = hnn.MobiusConcat if args.encoder_metric == cs.HY else hnn.EuclConcat
self.ctx_concat = ctx_concat_layer(ctx_dim * 2, ctx_dim)
self.ctx_attn = DistanceAttention(args, args.context_len * 2 + 2,
ctx_dim * 2)
# full concat of mention and context
input_classif_dim = men_dim + char_dim + ctx_dim * 2
full_concat_layer = hnn.MobiusConcat if args.concat_metric == cs.HY else hnn.EuclConcat
self.full_concat = full_concat_layer(input_classif_dim,
men_dim,
second_input_dim=ctx_dim * 2,
third_input_dim=char_dim)
# classifier
classifier_layer = hnn.MobiusMLR if args.mlr_metric == cs.HY else hnn.EuclMLR
self.classifier = classifier_layer(input_classif_dim,
vocabs[cs.TYPE_VOCAB].size())
self.attn_to_concat_map = define_mapping(args.attn_metric,
args.concat_metric, args.c)
self.concat_to_mlr_map = define_mapping(args.concat_metric,
args.mlr_metric, args.c)
def __init__(self,
*args,
alpha=None,
manifold=geoopt.Euclidean(),
**kwargs):
super().__init__(*args, **kwargs)
weight = self._parameters.pop("weight")
self._weight_shape = weight.shape
self.weight_orig = geoopt.ManifoldParameter(weight.data.reshape(
weight.shape[0], -1),
manifold=manifold)
with torch.no_grad():
self.weight_orig.proj_()
if alpha is None:
self.alpha_orig = nn.Parameter(torch.zeros(self._weight_shape[0]))
else:
self.alpha_orig = alpha
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 __init__(self, char_vocab, args):
super(MentionEncoder, self).__init__()
self.mention_output_dim = args.space_dims * 2
self.char_output_dim = args.space_dims
if args.encoder_metric == cs.HY:
self.manifold = gt.PoincareBall()
self.word2space = hnn.MobiusLinear(args.word_emb_size,
self.mention_output_dim,
hyperbolic_input=True,
hyperbolic_bias=True,
nonlin=get_nonlin(
args.men_nonlin))
self.non_lin = lambda x: x
self.char_rnn = hnn.MobiusRNN(args.space_dims, args.space_dims)
else:
self.manifold = gt.Euclidean()
self.word2space = nn.Linear(args.word_emb_size,
self.mention_output_dim)
self.non_lin = get_nonlin(args.men_nonlin)
self.char_rnn = hnn.EuclRNN(args.space_dims, args.space_dims)
self.input_dropout = nn.Dropout(p=args.input_dropout)
self.mention_attn = DistanceAttention(args, args.mention_len + 1,
self.mention_output_dim)
self.char_mapping = define_mapping(args.encoder_metric,
args.attn_metric, args.c)
self.char_midpoint = hnn.mobius_midpoint if args.attn_metric == cs.HY else hnn.euclidean_midpoint
# char embeds
with torch.no_grad():
char_embeds = init_embeddings(char_vocab.size(),
self.char_output_dim)
if args.encoder_metric == cs.HY:
char_embeds = pmath.expmap0(char_embeds, k=self.manifold.k)
self.char_lut = gt.ManifoldParameter(char_embeds,
manifold=self.manifold)
def test_fails_Euclidean():
with pytest.raises(ValueError):
manifold = geoopt.Euclidean(ndim=1)
manifold.random_normal(())
def test_random_R():
manifold = geoopt.Euclidean()
point = manifold.random_normal(10, 10)
manifold.assert_check_point_on_manifold(point)
assert point.manifold is manifold
def __init__(
self,
vocab_size,
embedding_dim,
hidden_dim,
project_dim,
cell_type="eucl_rnn",
embedding_type="eucl",
decision_type="eucl",
use_distance_as_feature=True,
device=None,
num_layers=1,
num_classes=1,
c=1.0,
):
super(RNNBase, self).__init__()
(cell_type, embedding_type,
decision_type) = map(str.lower,
[cell_type, embedding_type, decision_type])
if embedding_type == "eucl":
self.embedding = hyrnn.LookupEmbedding(vocab_size,
embedding_dim,
manifold=geoopt.Euclidean())
with torch.no_grad():
self.embedding.weight.normal_()
elif embedding_type == "hyp":
self.embedding = hyrnn.LookupEmbedding(
vocab_size,
embedding_dim,
manifold=geoopt.PoincareBall(c=c),
)
with torch.no_grad():
self.embedding.weight.set_(
pmath.expmap0(self.embedding.weight.normal_() / 10, c=c))
else:
raise NotImplementedError(
"Unsuported embedding type: {0}".format(embedding_type))
self.embedding_type = embedding_type
if decision_type == "eucl":
self.projector = nn.Linear(hidden_dim * 2, project_dim)
self.logits = nn.Linear(project_dim, num_classes)
elif decision_type == "hyp":
self.projector_source = hyrnn.MobiusLinear(hidden_dim,
project_dim,
c=c)
self.projector_target = hyrnn.MobiusLinear(hidden_dim,
project_dim,
c=c)
self.logits = hyrnn.MobiusDist2Hyperplane(project_dim, num_classes)
else:
raise NotImplementedError(
"Unsuported decision type: {0}".format(decision_type))
self.ball = geoopt.PoincareBall(c)
if use_distance_as_feature:
if decision_type == "eucl":
self.dist_bias = nn.Parameter(torch.zeros(project_dim))
else:
self.dist_bias = geoopt.ManifoldParameter(
torch.zeros(project_dim), manifold=self.ball)
else:
self.register_buffer("dist_bias", None)
self.decision_type = decision_type
self.use_distance_as_feature = use_distance_as_feature
self.device = device # declaring device here due to fact we are using catalyst
self.num_layers = num_layers
self.hidden_dim = hidden_dim
self.c = c
if cell_type == "eucl_rnn":
self.cell = nn.RNN
elif cell_type == "eucl_gru":
self.cell = nn.GRU
elif cell_type == "hyp_gru":
self.cell = functools.partial(hyrnn.MobiusGRU, c=c)
else:
raise NotImplementedError(
"Unsuported cell type: {0}".format(cell_type))
self.cell_type = cell_type
self.cell_source = self.cell(embedding_dim, self.hidden_dim,
self.num_layers)
self.cell_target = self.cell(embedding_dim, self.hidden_dim,
self.num_layers)
def test_fails_Euclidean():
with pytest.raises(ValueError):
manifold = geoopt.Euclidean(ndim=1)
manifold.origin(())
def test_random_Euclidean():
manifold = geoopt.Euclidean(ndim=1)
point = manifold.origin(10, 10)
manifold.assert_check_point_on_manifold(point)
assert point.manifold is manifold
def get_prod_hyeu_manifold(dims):
poincare = gt.PoincareBall()
euclidean = gt.Euclidean(1)
return gt.ProductManifold((poincare, dims // 2), (euclidean, dims // 2))
