Hyperbolic API

HeadPoints

Bases: NamedTuple

The points of any geometry head (Req 12) for a batch of B vectors in dimension d.

Attributes:

Name Type Description
tangent Tensor

The coordinates the export writes (B, d): the bounded tangent vector at the origin, r · û, under hyperbolic; v under Euclidean; û under spherical.
embedding Tensor

The arm's point: the Lorentz point exp_o(r · û) = (cosh r, sinh r · û) at c = 1 (B, d + 1) under hyperbolic; tangent itself (B, d) otherwise.
radius Tensor

The polar radius (B,): r, the point's geodesic distance from the origin, under hyperbolic; ‖v‖ under Euclidean; 1 under spherical, or 0 where v is 0.
direction Tensor

û, the unit direction of the head's input (B, d); zero where the input is zero.

HyperbolicHead

Bases: Module

The hyperbolic head (spec 4.2): a gradient-passing bound on the radius, then the exp map at the origin of the c = 1 hyperboloid. It has no parameters and no curvature.

From the projection's output v, with ν = ‖v‖ and the bound R, the radius is r = R · tanh(ν / R) and the direction is û = v / ν, so r ≤ R for every v. The radius takes gradient at any length (Req 13): dr/dν = sech²(ν / R) is positive, about 0.61 at a six-digit code's target r = 5 under R = 8. The interim cap at norm 2 passed about 1e-7 at its saturated points. In float32 the derivative rounds to 0 only from ν ≈ 8.7R, where tanh rounds to 1. At v = 0 the direction and the radius are 0, so the point is the origin, and every gradient there is finite.

It is Req 12's hyperbolic arm, the one arm with the radial term. geometry names the arm, distance the decoding distance for its points and radial whether the radial term applies; pair_distance is its training distance, and read_points maps exported tangents to the points its distance reads. The flat arms' heads are in text_model/heads.py.

Parameters:

Name Type Description Default
radius_bound float

R, the bound on every radius (model.radius_bound).

 8.0

Raises:

Type Description
ValueError

If radius_bound is not a positive finite number.

forward(vectors)

Bound each vector's radius and map it to the hyperboloid, in the vectors' dtype.

Parameters:

Name Type Description Default
vectors Tensor

The projection's outputs, (B, d).

 required

Returns:

Type Description
HeadPoints

The bounded tangent, the Lorentz point, the radius and the direction.

pair_distance(radius_a, direction_a, radius_b, direction_b) staticmethod

The training distance between two sets of the head's points: polar_distance.

read_points(tangent) staticmethod

Exported bounded tangents as Lorentz points, float64 on the CPU: exp_map_origin.

LorentzOps

Static utility class for Lorentz model operations. Provides functions for mapping between hyperboloid and tangent space, and computing distances.

exp_map_zero(x_tan, c=1.0) staticmethod

Exponential map from tangent space at origin to hyperboloid.

Maps a tangent vector at the origin to a point on the Lorentz hyperboloid. The output satisfies the Lorentz constraint: ||x_spatial||^2 - x0^2 = -1/c

Parameters:

Name Type Description Default
x_tan Tensor

Tangent vector, shape (batch_size, embedding_dim+1)

 required
c float

Curvature parameter (default: 1.0)

 1.0

Returns:

Type Description
Tensor

Point on hyperboloid, shape (batch_size, embedding_dim+1)

log_map_zero(x_hyp, c=1.0) staticmethod

Logarithmic map from hyperboloid to tangent space at origin.

Maps a point on the Lorentz hyperboloid to the tangent space at the origin. Inverse of exp_map_zero.

Parameters:

Name Type Description Default
x_hyp Tensor

Point on hyperboloid, shape (batch_size, embedding_dim+1) Must satisfy ||x_spatial||^2 - x0^2 = -1/c

 required
c float

Curvature parameter (default: 1.0)

 1.0

Returns:

Type Description
Tensor

Tangent vector, shape (batch_size, embedding_dim+1)

lorentz_distance(u, v, c=1.0) staticmethod

Compute Lorentzian distance between two points on the hyperboloid.

Parameters:

Name Type Description Default
u Tensor

First point on hyperboloid, shape (batch_size, embedding_dim+1)

 required
v Tensor

Second point on hyperboloid, shape (batch_size, embedding_dim+1)

 required
c float

Curvature parameter (default: 1.0)

 1.0

Returns:

Type Description
Tensor

Distances, shape (batch_size,)

check_lorentz_manifold_validity(embeddings, curvature=1.0, tolerance=0.001)

Check if embeddings satisfy the Lorentz hyperboloid constraint.

For valid points: -x₀² + x₁² + ... + xₙ² = -1/c

Parameters:

Name Type Description Default
embeddings Tensor

Hyperbolic embeddings of shape (batch_size, embedding_dim+1)

 required
curvature float

Curvature parameter c

 1.0
tolerance float

Tolerance for constraint violation

 0.001

Returns:

Type Description
Tuple[bool, Tensor, Tensor]

Tuple of: - is_valid: Boolean indicating if all points are valid - lorentz_norms: Lorentz inner product for each point (should be -1/c) - violations: Magnitude of constraint violations

compute_hyperbolic_radii(embeddings)

Extract hyperbolic radii (time coordinates) from Lorentz embeddings.

The time coordinate x₀ represents the hyperbolic radius (distance from origin).

Parameters:

Name Type Description Default
embeddings Tensor

Hyperbolic embeddings of shape (batch_size, embedding_dim+1)

 required

Returns:

Type Description
Tensor

Hyperbolic radii of shape (batch_size,)

exp_map_origin(tangent)

The exponential map at the origin of the curvature -1 hyperboloid (c = 1), in float64.

It is HyperbolicHead's map, computed in float64 on the CPU whatever the tangent's device and dtype: the hyperbolic arm's read map (HyperbolicHead.read_points).

Parameters:

Name Type Description Default
tangent Tensor

Tangent vectors at the origin (N, d).

 required

Returns:

Type Description
Tensor

(time, space) rows (N, d + 1), float64 on the CPU.

polar_distance(radius_a, direction_a, radius_b, direction_b)

The geodesic distance at c = 1 between every point of one set and every point of another, from each point's radius r and direction û (spec 4.2).

The distance is the hyperbolic law of cosines in half-angle form, so it equals arcosh(−⟨x, y⟩_L) exactly:

d(x, y) = 2 · asinh(√(sinh²((r_x − r_y) / 2) + sinh r_x · sinh r_y · ‖û_x − û_y‖² / 4))

Every term is non-negative, so nothing cancels, and float32 resolves it at every radius the head's bound allows. ‖û_x − û_y‖² comes from explicit differences, not from 2 − 2 û_x · û_y, which cancels at small angles. The ops are elementwise, so under autocast float32 inputs give float32 distances. The square root is guarded at zero: a zero separation has distance 0, and the value and its gradient stay finite.

Parameters:

Name Type Description Default
radius_a Tensor

The first set's radii, (A,).

 required
direction_a Tensor

Its directions, (A, d): unit vectors, or 0 at the origin.

 required
radius_b Tensor

The second set's radii, (B,).

 required
direction_b Tensor

Its directions, (B, d).

 required

Returns:

Type Description
Tensor

The distances, (A, B), in the inputs' dtype.

Raises:

Type Description
ValueError

If the shapes are not (A,), (A, d), (B,) and (B, d).