Skip to content

Functional

Stateless operations. No learnable parameters.

Activations

GeometricGELU

Bases: CliffordModule

Geometric GELU activation: x' = x * GELU(||x|| + b) / ||x||.

Scales magnitude while preserving direction.

Attributes:

Name Type Description
algebra CliffordAlgebra

The algebra instance.
bias Parameter

Learnable bias added to norm.
Source code in functional/activation.py 
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
class GeometricGELU(CliffordModule):
    """Geometric GELU activation: x' = x * GELU(||x|| + b) / ||x||.

    Scales magnitude while preserving direction.

    Attributes:
        algebra (CliffordAlgebra): The algebra instance.
        bias (torch.nn.Parameter): Learnable bias added to norm.
    """

    def __init__(self, algebra, channels: int = 1):
        """Initialize Geometric GELU.

        Args:
            algebra (CliffordAlgebra): The algebra instance.
            channels (int): Number of channels.
        """
        super().__init__(algebra)
        self.bias = nn.Parameter(torch.zeros(channels))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Apply geometric GELU activation.

        Args:
            x (torch.Tensor): Input multivector [..., Dim].

        Returns:
            torch.Tensor: Activated multivector.
        """
        norm = x.norm(dim=-1, keepdim=True)

        eps = 1e-6
        scale = F.gelu(norm + self.bias.view(1, -1, 1)) / (norm + eps)

        return x * scale

__init__(algebra, channels=1)

Initialize Geometric GELU.

Parameters:

Name Type Description Default
algebra CliffordAlgebra

The algebra instance.

 required
channels int

Number of channels.

 1
Source code in functional/activation.py 
29
30
31
32
33
34
35
36
37
def __init__(self, algebra, channels: int = 1):
    """Initialize Geometric GELU.

    Args:
        algebra (CliffordAlgebra): The algebra instance.
        channels (int): Number of channels.
    """
    super().__init__(algebra)
    self.bias = nn.Parameter(torch.zeros(channels))

forward(x)

Apply geometric GELU activation.

Parameters:

Name Type Description Default
x Tensor

Input multivector [..., Dim].

 required

Returns:

Type Description
Tensor

torch.Tensor: Activated multivector.
Source code in functional/activation.py 
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Apply geometric GELU activation.

    Args:
        x (torch.Tensor): Input multivector [..., Dim].

    Returns:
        torch.Tensor: Activated multivector.
    """
    norm = x.norm(dim=-1, keepdim=True)

    eps = 1e-6
    scale = F.gelu(norm + self.bias.view(1, -1, 1)) / (norm + eps)

    return x * scale

GradeSwish

Bases: CliffordModule

Per-grade gated activation.

Each grade receives an independent sigmoid gate based on its norm.

Attributes:

Name Type Description
algebra CliffordAlgebra

The algebra instance.
n_grades int

Number of grades.
grade_weights Parameter

Weights for each grade gate.
grade_biases Parameter

Biases for each grade gate.
Source code in functional/activation.py 
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
class GradeSwish(CliffordModule):
    """Per-grade gated activation.

    Each grade receives an independent sigmoid gate based on its norm.

    Attributes:
        algebra (CliffordAlgebra): The algebra instance.
        n_grades (int): Number of grades.
        grade_weights (torch.nn.Parameter): Weights for each grade gate.
        grade_biases (torch.nn.Parameter): Biases for each grade gate.
    """

    def __init__(self, algebra, channels: int = 1):
        """Initialize Grade Swish.

        Args:
            algebra (CliffordAlgebra): The algebra instance.
            channels (int): Number of channels.
        """
        super().__init__(algebra)
        self.n_grades = self.algebra.n + 1

        self.grade_weights = nn.Parameter(torch.ones(self.n_grades))
        self.grade_biases = nn.Parameter(torch.zeros(self.n_grades))

        # Reuse algebra's precomputed grade_index instead of building our own
        self.register_buffer('_grade_index', self.algebra.grade_index.clone())

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Apply per-grade gating.

        Args:
            x (torch.Tensor): Input multivector [..., Dim].

        Returns:
            torch.Tensor: Activated multivector.
        """
        D = self.algebra.dim
        G = self.n_grades

        # Square, scatter-add by grade, sqrt -> per-grade norms
        x_sq = x * x  # [..., D]
        batch_shape = x.shape[:-1]
        grade_idx = self._grade_index.expand(*batch_shape, D)  # [..., D]

        norm_sq = torch.zeros(*batch_shape, G, device=x.device, dtype=x.dtype)
        norm_sq.scatter_add_(-1, grade_idx, x_sq)  # [..., G]
        norms = torch.sqrt(norm_sq.clamp(min=1e-12))  # [..., G]

        # Compute gates: sigmoid(w * norm + b) for each grade
        gates = torch.sigmoid(
            self.grade_weights * norms + self.grade_biases
        )  # [..., G]

        # Broadcast gate per component: lookup gate[grade_map[d]] for each d
        per_component_gate = gates.gather(-1, grade_idx)  # [..., D]

        return x * per_component_gate

__init__(algebra, channels=1)

Initialize Grade Swish.

Parameters:

Name Type Description Default
algebra CliffordAlgebra

The algebra instance.

 required
channels int

Number of channels.

 1
Source code in functional/activation.py 
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
def __init__(self, algebra, channels: int = 1):
    """Initialize Grade Swish.

    Args:
        algebra (CliffordAlgebra): The algebra instance.
        channels (int): Number of channels.
    """
    super().__init__(algebra)
    self.n_grades = self.algebra.n + 1

    self.grade_weights = nn.Parameter(torch.ones(self.n_grades))
    self.grade_biases = nn.Parameter(torch.zeros(self.n_grades))

    # Reuse algebra's precomputed grade_index instead of building our own
    self.register_buffer('_grade_index', self.algebra.grade_index.clone())

forward(x)

Apply per-grade gating.

Parameters:

Name Type Description Default
x Tensor

Input multivector [..., Dim].

 required

Returns:

Type Description
Tensor

torch.Tensor: Activated multivector.
Source code in functional/activation.py 
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Apply per-grade gating.

    Args:
        x (torch.Tensor): Input multivector [..., Dim].

    Returns:
        torch.Tensor: Activated multivector.
    """
    D = self.algebra.dim
    G = self.n_grades

    # Square, scatter-add by grade, sqrt -> per-grade norms
    x_sq = x * x  # [..., D]
    batch_shape = x.shape[:-1]
    grade_idx = self._grade_index.expand(*batch_shape, D)  # [..., D]

    norm_sq = torch.zeros(*batch_shape, G, device=x.device, dtype=x.dtype)
    norm_sq.scatter_add_(-1, grade_idx, x_sq)  # [..., G]
    norms = torch.sqrt(norm_sq.clamp(min=1e-12))  # [..., G]

    # Compute gates: sigmoid(w * norm + b) for each grade
    gates = torch.sigmoid(
        self.grade_weights * norms + self.grade_biases
    )  # [..., G]

    # Broadcast gate per component: lookup gate[grade_map[d]] for each d
    per_component_gate = gates.gather(-1, grade_idx)  # [..., D]

    return x * per_component_gate

Losses

ChamferDistance

Bases: Module

Symmetric Chamfer distance between two point clouds.

CD(P, Q) = (1/|P|) sum_p min_q ||p-q||^2 + (1/|Q|) sum_q min_p ||q-p||^2

Standard metric for 3D point cloud reconstruction and generation.

Source code in functional/loss.py 
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
class ChamferDistance(nn.Module):
    """Symmetric Chamfer distance between two point clouds.

    CD(P, Q) = (1/|P|) sum_p min_q ||p-q||^2 + (1/|Q|) sum_q min_p ||q-p||^2

    Standard metric for 3D point cloud reconstruction and generation.
    """

    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        """Compute Chamfer distance.

        Args:
            pred: Predicted point cloud [B, M, 3].
            target: Target point cloud [B, N, 3].

        Returns:
            Chamfer distance (scalar).
        """
        diff = pred.unsqueeze(2) - target.unsqueeze(1)
        dist_sq = (diff ** 2).sum(dim=-1)
        min_dist_pred = dist_sq.min(dim=2)[0].mean(dim=1)
        min_dist_target = dist_sq.min(dim=1)[0].mean(dim=1)
        return (min_dist_pred + min_dist_target).mean()

forward(pred, target)

Compute Chamfer distance.

Parameters:

Name Type Description Default
pred Tensor

Predicted point cloud [B, M, 3].

 required
target Tensor

Target point cloud [B, N, 3].

 required

Returns:

Type Description
Tensor

Chamfer distance (scalar).
Source code in functional/loss.py 
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    """Compute Chamfer distance.

    Args:
        pred: Predicted point cloud [B, M, 3].
        target: Target point cloud [B, N, 3].

    Returns:
        Chamfer distance (scalar).
    """
    diff = pred.unsqueeze(2) - target.unsqueeze(1)
    dist_sq = (diff ** 2).sum(dim=-1)
    min_dist_pred = dist_sq.min(dim=2)[0].mean(dim=1)
    min_dist_target = dist_sq.min(dim=1)[0].mean(dim=1)
    return (min_dist_pred + min_dist_target).mean()

ConservativeLoss

Bases: Module

Enforces F = -grad(E) conservative force constraint.

Physics: forces should be the negative gradient of energy with respect to atomic positions. Used in molecular dynamics tasks.

Source code in functional/loss.py 
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
class ConservativeLoss(nn.Module):
    """Enforces F = -grad(E) conservative force constraint.

    Physics: forces should be the negative gradient of energy
    with respect to atomic positions. Used in molecular dynamics tasks.
    """

    def forward(self, energy: torch.Tensor, force_pred: torch.Tensor,
                pos: torch.Tensor) -> torch.Tensor:
        """Compute conservative force loss.

        Args:
            energy: Predicted energy (scalar, requires grad graph).
            force_pred: Predicted forces [N, 3].
            pos: Atom positions [N, 3] (must have requires_grad=True).

        Returns:
            MSE between predicted forces and -grad(E).
        """
        force_from_energy = -torch.autograd.grad(
            energy.sum(), pos,
            create_graph=True, retain_graph=True
        )[0]
        return F.mse_loss(force_pred, force_from_energy)

forward(energy, force_pred, pos)

Compute conservative force loss.

Parameters:

Name Type Description Default
energy Tensor

Predicted energy (scalar, requires grad graph).

 required
force_pred Tensor

Predicted forces [N, 3].

 required
pos Tensor

Atom positions [N, 3] (must have requires_grad=True).

 required

Returns:

Type Description
Tensor

MSE between predicted forces and -grad(E).
Source code in functional/loss.py 
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
def forward(self, energy: torch.Tensor, force_pred: torch.Tensor,
            pos: torch.Tensor) -> torch.Tensor:
    """Compute conservative force loss.

    Args:
        energy: Predicted energy (scalar, requires grad graph).
        force_pred: Predicted forces [N, 3].
        pos: Atom positions [N, 3] (must have requires_grad=True).

    Returns:
        MSE between predicted forces and -grad(E).
    """
    force_from_energy = -torch.autograd.grad(
        energy.sum(), pos,
        create_graph=True, retain_graph=True
    )[0]
    return F.mse_loss(force_pred, force_from_energy)

PhysicsInformedLoss

Bases: Module

Physics-informed loss combining MSE with conservation penalty.

Enforces that global weighted mean of each variable is approximately conserved between forecast and target. Used in weather forecasting.

Source code in functional/loss.py 
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
class PhysicsInformedLoss(nn.Module):
    """Physics-informed loss combining MSE with conservation penalty.

    Enforces that global weighted mean of each variable is approximately
    conserved between forecast and target. Used in weather forecasting.
    """

    def __init__(self, physics_weight: float = 0.1):
        super().__init__()
        self.physics_weight = physics_weight

    def forward(self, forecast: torch.Tensor, target: torch.Tensor,
                lat_weights: torch.Tensor = None) -> torch.Tensor:
        """Compute physics-informed loss.

        Args:
            forecast: Predicted state [B, H, W, C].
            target: Target state [B, H, W, C].
            lat_weights: Latitude area weights [H].

        Returns:
            Combined MSE + conservation penalty.
        """
        mse_loss = F.mse_loss(forecast, target)

        if lat_weights is not None and forecast.dim() == 4:
            w = lat_weights.view(1, -1, 1, 1).to(forecast.device)
            forecast_mean = (forecast * w).sum(dim=[1, 2]) / w.sum()
            target_mean = (target * w).sum(dim=[1, 2]) / w.sum()
        else:
            forecast_mean = forecast.mean(dim=list(range(1, forecast.dim() - 1)))
            target_mean = target.mean(dim=list(range(1, target.dim() - 1)))

        conservation_loss = F.mse_loss(forecast_mean, target_mean)
        return mse_loss + self.physics_weight * conservation_loss

forward(forecast, target, lat_weights=None)

Compute physics-informed loss.

Parameters:

Name Type Description Default
forecast Tensor

Predicted state [B, H, W, C].

 required
target Tensor

Target state [B, H, W, C].

 required
lat_weights Tensor

Latitude area weights [H].

 None

Returns:

Type Description
Tensor

Combined MSE + conservation penalty.
Source code in functional/loss.py 
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
def forward(self, forecast: torch.Tensor, target: torch.Tensor,
            lat_weights: torch.Tensor = None) -> torch.Tensor:
    """Compute physics-informed loss.

    Args:
        forecast: Predicted state [B, H, W, C].
        target: Target state [B, H, W, C].
        lat_weights: Latitude area weights [H].

    Returns:
        Combined MSE + conservation penalty.
    """
    mse_loss = F.mse_loss(forecast, target)

    if lat_weights is not None and forecast.dim() == 4:
        w = lat_weights.view(1, -1, 1, 1).to(forecast.device)
        forecast_mean = (forecast * w).sum(dim=[1, 2]) / w.sum()
        target_mean = (target * w).sum(dim=[1, 2]) / w.sum()
    else:
        forecast_mean = forecast.mean(dim=list(range(1, forecast.dim() - 1)))
        target_mean = target.mean(dim=list(range(1, target.dim() - 1)))

    conservation_loss = F.mse_loss(forecast_mean, target_mean)
    return mse_loss + self.physics_weight * conservation_loss

GeometricMSELoss

Bases: CliffordModule

Geometric MSE. Euclidean distance in embedding space.

Standard MSE on coefficients.

Source code in functional/loss.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
class GeometricMSELoss(CliffordModule):
    """Geometric MSE. Euclidean distance in embedding space.

    Standard MSE on coefficients.
    """

    def __init__(self, algebra):
        """Initialize the geometric MSE loss."""
        super().__init__(algebra)

    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        """MSE."""
        return F.mse_loss(pred, target, reduction='mean')

__init__(algebra)

Initialize the geometric MSE loss.

Source code in functional/loss.py 
20
21
22
def __init__(self, algebra):
    """Initialize the geometric MSE loss."""
    super().__init__(algebra)

forward(pred, target)

MSE.

Source code in functional/loss.py 
24
25
26
def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    """MSE."""
    return F.mse_loss(pred, target, reduction='mean')

SubspaceLoss

Bases: CliffordModule

Subspace Loss. Enforces grade constraints.

Penalizes energy in forbidden grades.

Source code in functional/loss.py 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
class SubspaceLoss(CliffordModule):
    """Subspace Loss. Enforces grade constraints.

    Penalizes energy in forbidden grades.
    """

    def __init__(self, algebra, target_indices: list = None, exclude_indices: list = None):
        """Initialize grade constraint penalties."""
        super().__init__(algebra)

        if target_indices is not None:
            mask = torch.ones(self.algebra.dim, dtype=torch.bool)
            mask[target_indices] = False
        elif exclude_indices is not None:
            mask = torch.zeros(self.algebra.dim, dtype=torch.bool)
            mask[exclude_indices] = True
        else:
            raise ValueError("Must provide target_indices or exclude_indices")

        self.register_buffer('penalty_mask', mask)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Penalizes deviations."""
        penalty_components = x[..., self.penalty_mask]
        loss = (penalty_components ** 2).sum(dim=-1).mean()
        return loss

__init__(algebra, target_indices=None, exclude_indices=None)

Initialize grade constraint penalties.

Source code in functional/loss.py 
34
35
36
37
38
39
40
41
42
43
44
45
46
47
def __init__(self, algebra, target_indices: list = None, exclude_indices: list = None):
    """Initialize grade constraint penalties."""
    super().__init__(algebra)

    if target_indices is not None:
        mask = torch.ones(self.algebra.dim, dtype=torch.bool)
        mask[target_indices] = False
    elif exclude_indices is not None:
        mask = torch.zeros(self.algebra.dim, dtype=torch.bool)
        mask[exclude_indices] = True
    else:
        raise ValueError("Must provide target_indices or exclude_indices")

    self.register_buffer('penalty_mask', mask)

forward(x)

Penalizes deviations.

Source code in functional/loss.py 
49
50
51
52
53
def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Penalizes deviations."""
    penalty_components = x[..., self.penalty_mask]
    loss = (penalty_components ** 2).sum(dim=-1).mean()
    return loss

IsometryLoss

Bases: CliffordModule

Isometry loss enforcing metric norm preservation.

Ensures transformations preserve the metric norm.

Source code in functional/loss.py 
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
class IsometryLoss(CliffordModule):
    """Isometry loss enforcing metric norm preservation.

    Ensures transformations preserve the metric norm.
    """

    def __init__(self, algebra):
        """Initialize isometry loss with metric diagonal."""
        super().__init__(algebra)
        self.register_buffer('metric_diag', self._compute_metric_diagonal())

    def _compute_metric_diagonal(self):
        """Finds the signature."""
        basis = torch.eye(self.algebra.dim, device=self.algebra.device)
        sq = self.algebra.geometric_product(basis, basis)
        diag = sq[:, 0]
        return diag

    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        """Compares norms."""
        metric_diag = self.metric_diag
        pred_sq = (pred ** 2) * metric_diag
        target_sq = (target ** 2) * metric_diag

        pred_norm = pred_sq.sum(dim=-1)
        target_norm = target_sq.sum(dim=-1)

        return F.mse_loss(pred_norm, target_norm)

__init__(algebra)

Initialize isometry loss with metric diagonal.

Source code in functional/loss.py 
61
62
63
64
def __init__(self, algebra):
    """Initialize isometry loss with metric diagonal."""
    super().__init__(algebra)
    self.register_buffer('metric_diag', self._compute_metric_diagonal())

forward(pred, target)

Compares norms.

Source code in functional/loss.py 
73
74
75
76
77
78
79
80
81
82
def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    """Compares norms."""
    metric_diag = self.metric_diag
    pred_sq = (pred ** 2) * metric_diag
    target_sq = (target ** 2) * metric_diag

    pred_norm = pred_sq.sum(dim=-1)
    target_norm = target_sq.sum(dim=-1)

    return F.mse_loss(pred_norm, target_norm)

BivectorRegularization

Bases: CliffordModule

Bivector regularization enforcing grade-2 purity.

Penalizes energy outside the target grade (default: grade 2).

Source code in functional/loss.py 
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
class BivectorRegularization(CliffordModule):
    """Bivector regularization enforcing grade-2 purity.

    Penalizes energy outside the target grade (default: grade 2).
    """

    def __init__(self, algebra, grade=2):
        """Initialize bivector regularization."""
        super().__init__(algebra)
        self.grade = grade

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Penalizes non-bivector parts."""
        target_part = self.algebra.grade_projection(x, self.grade)
        residual = x - target_part
        return (residual ** 2).sum(dim=-1).mean()

__init__(algebra, grade=2)

Initialize bivector regularization.

Source code in functional/loss.py 
90
91
92
93
def __init__(self, algebra, grade=2):
    """Initialize bivector regularization."""
    super().__init__(algebra)
    self.grade = grade

forward(x)

Penalizes non-bivector parts.

Source code in functional/loss.py 
95
96
97
98
99
def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Penalizes non-bivector parts."""
    target_part = self.algebra.grade_projection(x, self.grade)
    residual = x - target_part
    return (residual ** 2).sum(dim=-1).mean()