src/manifold/models/layers/__init__.py

from manifold.models.layers.attention import (
    RotaryPositionEncoding,
    MultiHeadLinearAttention,
    linear_attention,
    causal_linear_attention,
    apply_rotary_pos_emb,
)
from manifold.models.layers.hyperbolic import (
    expmap0,
    logmap0,
    hyperbolic_distance_tangent,
    poincare_distance,
    project_to_ball,
    TangentSpaceProjection,
    HyperbolicMLP,
    HyperbolicDistanceLayer,
)
from manifold.models.layers.physics import (
    PhysicsConstraints,
    PhysicsConstraintLayer,
    SignalDependentNoiseChecker,
    compute_jerk,
    compute_jerk_violation,
    compute_fitts_violation,
    compute_reaction_time_violation,
)

__all__ = [
    "RotaryPositionEncoding",
    "MultiHeadLinearAttention",
    "linear_attention",
    "causal_linear_attention",
    "apply_rotary_pos_emb",
    "expmap0",
    "logmap0",
    "hyperbolic_distance_tangent",
    "poincare_distance",
    "project_to_ball",
    "TangentSpaceProjection",
    "HyperbolicMLP",
    "HyperbolicDistanceLayer",
    "PhysicsConstraints",
    "PhysicsConstraintLayer",
    "SignalDependentNoiseChecker",
    "compute_jerk",
    "compute_jerk_violation",
    "compute_fitts_violation",
    "compute_reaction_time_violation",
]