def _check_point_on_manifold(self, x, *, atol=1e-5, rtol=1e-5):
px = math.project(x, c=self.c)
ok = torch.allclose(x, px, atol=atol, rtol=rtol)
if not ok:
reason = "'x' norm lies out of the bounds [-1/sqrt(c)+eps, 1/sqrt(c)-eps]"
else:
reason = None
return ok, reason
def expmap(self, x, u, *, project=True, dim=-1):
res = math.expmap(x, u, c=self.c, dim=dim)
if project:
return math.project(res, c=self.c, dim=dim)
else:
return res
def projx(self, x, dim=-1):
return math.project(x, c=self.c, dim=dim)
def retr(self, x, u, *, dim=-1):
# always assume u is scaled properly
approx = x + u
return math.project(approx, c=self.c, dim=dim)
def mobius_fn_apply_chain(self, x, *fns, project=True, dim=-1):
res = math.mobius_fn_apply_chain(x, *fns, c=self.c, dim=dim)
if project:
return math.project(res, c=self.c, dim=dim)
else:
return res
def expmap0(self, u, *, dim=-1, project=True):
res = math.expmap0(u, c=self.c, dim=dim)
if project:
return math.project(res, c=self.c, dim=dim)
else:
return res
def mobius_fn_apply(self, fn, x, *args, dim=-1, project=True, **kwargs):
res = math.mobius_fn_apply(fn, x, *args, c=self.c, dim=dim, **kwargs)
if project:
return math.project(res, c=self.c, dim=dim)
else:
return res
def geodesic_unit(self, t, x, u, *, dim=-1, project=True):
res = math.geodesic_unit(t, x, u, c=self.c, dim=dim)
if project:
return math.project(res, c=self.c, dim=dim)
else:
return res
def mobius_matvec(self, m, x, *, dim=-1, project=True):
res = math.mobius_matvec(m, x, c=self.c, dim=dim)
if project:
return math.project(res, c=self.c, dim=dim)
else:
return res
def mobius_pointwise_mul(self, w, x, *, dim=-1, project=True):
res = math.mobius_pointwise_mul(w, x, c=self.c, dim=dim)
if project:
return math.project(res, c=self.c, dim=dim)
else:
return res
def mobius_scalar_mul(self, r, x, *, dim=-1, project=True):
res = math.mobius_scalar_mul(r, x, c=self.c, dim=dim)
if project:
return math.project(res, c=self.c, dim=dim)
else:
return res
def mobius_cosub(self, x, y, *, dim=-1, project=True):
res = math.mobius_coadd(x, y, c=self.c, dim=dim)
if project:
return math.project(res, c=self.c, dim=dim)
else:
return res