Diffusion and flow matching

Stable SO(3) exponential and logarithm maps

HardRaw PyTorchAbout 30 min3 public examples plus private tests

FrameFlowFrameDiffRFdiffusionBioEmu

Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.

Your task

Implement the rotation-vector exponential map and the principal rotation-matrix logarithm independently. Handle both zero angle and rotations near pi using the stated stable branches.

Function signature

so3_exp_log(
    omega: torch.Tensor,
    R: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]

Worked example

Two independent quarter-turns

Use $\omega=[[0,0,\pi/2]]$ and independently supply the x-axis quarter-turn R=[[[1,0,0],[0,0,-1],[0,1,0]]]. The exponential returns the z-axis quarter-turn

\mathrm{exp\_R}=\begin{bmatrix}0&-1&0\\1&0&0\\0&0&1\end{bmatrix}

with its batch axis retained. For the supplied R, the antisymmetric part identifies the positive x axis and $(\operatorname{tr}R-1)/2=0$ , so the principal angle is $\pi/2$ . Return $\mathrm{log\_omega}=[[\pi/2,0,0]]$ , not the original omega. The two return shapes are (1,3,3) and (1,3).

Parameters

Parameter	Shape and meaning
`omega`torch.Tensor	$[\ldots, 3]$ Rotation vectors in radians
`R`torch.Tensor	$[\ldots, 3, 3]$ Independent proper rotation matrices

Return

tuple[torch.Tensor, torch.Tensor]

Return these tensors in the stated order.

Required behavior

Implement BOTH maps independently in one call: omega $[\ldots,3]$ produces exp_R $[\ldots,3,3]$ , and $R [\ldots,3,3]$ produces log_omega $[\ldots,3]$ . Their leading shapes match, but $R$ is not assumed to equal $\operatorname{Exp}(\omega)$ . The leading shape may be empty, one axis or multiple axes, all nonempty where present.
Rotations act on column vectors and map local coordinates to the global frame.
For $v=(x,y,z)$ , $\operatorname{hat}(v)=\begin{pmatrix}0&-z&y\\z&0&-x\\-y&x&0\end{pmatrix}$ . $\operatorname{Exp}(v)=I+A\operatorname{hat}(v)+B\operatorname{hat}(v)^2$ , where $q=v\cdot v$ , $A=\frac{\sin\sqrt q}{\sqrt q}$ , $B=\frac{1-\cos\sqrt q}{q}$ .
For $q<10^{-4}$ use $A=1-\frac q6+\frac{q^2}{120}-\frac{q^3}{5040}$ and $B=\frac12-\frac q{24}+\frac{q^2}{720}-\frac{q^3}{40320}$ ; at zero $\operatorname{Exp}(0)=I$ .
For $\operatorname{Log}(R)$ , set $v=\frac12\operatorname{vee}(R-R^\top)$ , $q=v\cdot v$ , $s=\sqrt{q}$ , $c=\frac{\operatorname{tr}R-1}{2}$ , $\theta=\operatorname{atan2}(s,c)$ . Ordinarily $\operatorname{Log}(R)=\frac{\theta}{s}v$ . If $q<10^{-8}$ and $c>0$ use $v\left(1+\frac q6+\frac{3q^2}{40}+\frac{5q^3}{112}\right)$ .
If $q<10^{-6}$ and $c<0$ , recover an axis from $P=\frac{R+R^\top-2cI}{2(1-c)}$ : choose the largest diagonal $(\mathrm{ties} x, \mathrm{then} y, \mathrm{then} z)$ , take its positive square root, and divide its row by that component; normalize the axis.
Orient it to have nonnegative dot with $v$ when $q>10^{-20}$ ; otherwise keep the chosen component positive. Return $\theta\cdot \mathrm{axis}$ .
For dormant branches clamp $1-c$ below by $10^{-6}$ and diagonal square-root arguments below by $10^{-12}$ ; these floors do not alter the selected near- $\pi$ axis. Evaluate every inactive square root/division safely before torch.where. $R$ must be a proper $\operatorname{SO}(3)$ matrix; no projection/repair is requested.
At $\pi$ axis sign is inherently ambiguous: $\operatorname{exp}(\operatorname{log}(R))$ is the geometric invariant. No global continuous logarithm exists at the $\pi$ cut locus. No masks, randomness, trainable weights or options are used. Floating inputs share dtype/device.

Why this matters

A rotation vector points along the rotation axis; its length is the angle in radians. The matrix describes the same action on column vectors. The two inputs in this task are independent.

Starter code

You write the body of this function in the BioTorch workspace. Your code runs in an isolated PyTorch sandbox; Run checks the public examples and Submit grades the full test suite, including gradients.

import torch


def so3_exp_log(
    omega: torch.Tensor,  # (..., 3)
    R: torch.Tensor,  # (..., 3, 3)
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return tuple[torch.Tensor, torch.Tensor]:
    
    exp_R: (..., 3, 3).
    log_omega: (..., 3)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement so3_exp_log")

Ready to implement it?

3 progressive hints, a solution walkthrough, and grading that checks values, shapes and gradients are waiting in the workspace.

Write and run your solution Sign in

Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.

Sources and further reading

A primitive exercise, not an end-to-end reproduction of the named models. No affiliation or endorsement is implied.