pymomentum.backend

PyMomentum Backend

High-performance implementations for forward kinematics and linear blend skinning operations.

pymomentum.backend.skel_state_backend module

Skeleton State Backend for PyMomentum

This module provides efficient forward kinematics and skinning operations using the skeleton state representation (8 parameters per joint: translation, quaternion, scale).

The skeleton state representation is more compact than the TRS backend and uses quaternions for rotation, making it suitable for applications requiring smooth interpolation and fewer parameters.

Performance Notes: This backend matches the behavior of the C++ Momentum code, ensuring consistency with the reference implementation. However, the TRS backend may be 25-50% faster in PyTorch due to not requiring quaternion normalization operations, though the exact performance difference may vary depending on the specific use case.

Key Functions: - global_skel_state_from_local_skel_state: Forward kinematics from local to global joint states - skin_points_from_skel_state: Linear blend skinning using skeleton states - local_skel_state_from_joint_params: Convert joint parameters to local states

Related Modules: - trs_backend: Alternative backend using separate translation/rotation/scale tensors - skel_state: Core skeleton state operations and utilities - quaternion: Quaternion math operations used by this backend

class pymomentum.backend.skel_state_backend.GlobalSkelStateFromLocalSkelStateJIT(*args, **kwargs)

Bases: Function

PyTorch autograd function for differentiable forward kinematics using skeleton states.

This class implements automatic differentiation for the forward kinematics operation, allowing gradients to flow from global joint states back to local joint states. The forward pass computes global states efficiently while saving intermediate results for the backward pass.

Note:: This class is used internally by global_skel_state_from_local_skel_state when not in JIT mode. It provides gradient computation capabilities that are not available in pure JIT-compiled functions.

static backward(ctx, grad_global_skel_state: Tensor, _0) → Tuple[Tensor, None]

Define a formula for differentiating the operation with backward mode automatic differentiation.

This function is to be overridden by all subclasses. (Defining this function is equivalent to defining the vjp function.)

It must accept a context ctx as the first argument, followed by as many outputs as the forward() returned (None will be passed in for non tensor outputs of the forward function), and it should return as many tensors, as there were inputs to forward(). Each argument is the gradient w.r.t the given output, and each returned value should be the gradient w.r.t. the corresponding input. If an input is not a Tensor or is a Tensor not requiring grads, you can just pass None as a gradient for that input.

The context can be used to retrieve tensors saved during the forward pass. It also has an attribute ctx.needs_input_grad as a tuple of booleans representing whether each input needs gradient. E.g., backward() will have ctx.needs_input_grad[0] = True if the first input to forward() needs gradient computed w.r.t. the output.

static forward(local_skel_state: Tensor, prefix_mul_indices: List[Tensor]) → Tuple[Tensor, List[Tensor]]

Define the forward of the custom autograd Function.

This function is to be overridden by all subclasses. There are two ways to define forward:

Usage 1 (Combined forward and ctx):

@staticmethod
def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any:
    pass

It must accept a context ctx as the first argument, followed by any number of arguments (tensors or other types).
See combining-forward-context for more details

Usage 2 (Separate forward and ctx):

@staticmethod
def forward(*args: Any, **kwargs: Any) -> Any:
    pass

@staticmethod
def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None:
    pass

The forward no longer accepts a ctx argument.
Instead, you must also override the torch.autograd.Function.setup_context() staticmethod to handle setting up the ctx object. output is the output of the forward, inputs are a Tuple of inputs to the forward.
See extending-autograd for more details

The context can be used to store arbitrary data that can be then retrieved during the backward pass. Tensors should not be stored directly on ctx (though this is not currently enforced for backward compatibility). Instead, tensors should be saved either with ctx.save_for_backward() if they are intended to be used in backward (equivalently, vjp) or ctx.save_for_forward() if they are intended to be used for in jvp.

static setup_context(ctx, inputs, outputs) → None

There are two ways to define the forward pass of an autograd.Function.

Either:

Override forward with the signature forward(ctx, *args, **kwargs). setup_context is not overridden. Setting up the ctx for backward happens inside the forward.
Override forward with the signature forward(*args, **kwargs) and override setup_context. Setting up the ctx for backward happens inside setup_context (as opposed to inside the forward)

See torch.autograd.Function.forward() and extending-autograd for more details.

pymomentum.backend.skel_state_backend.global_skel_state_from_local_skel_state(local_skel_state: Tensor, prefix_mul_indices: List[Tensor]) → Tensor

Compute global skeleton state from local joint transformations (user-facing wrapper).

This is the main entry point for forward kinematics using skeleton states. It automatically selects between JIT-compiled and autograd-enabled implementations based on the execution context.

Args:

local_skel_state: Local joint transformations, shape (batch_size, num_joints, 8).: Each joint contains [tx, ty, tz, qx, qy, qz, qw, s] parameters.
prefix_mul_indices: List of [child_index, parent_index] tensor pairs defining: the kinematic hierarchy traversal order.

Returns:

global_skel_state: Global joint transformations, shape (batch_size, num_joints, 8).: Each joint contains the composed transformation from root to that joint.

Note:

When called within torch.jit.script or torch.jit.trace context, uses the JIT-compiled implementation for maximum performance. Otherwise, uses the autograd-enabled version for gradient computation.

See Also:

global_skel_state_from_local_skel_state_impl(): JIT implementation local_skel_state_from_joint_params(): Convert joint parameters to local states

pymomentum.backend.skel_state_backend.global_skel_state_from_local_skel_state_backprop(global_skel_state: Tensor, grad_global_skel_state: Tensor, prefix_mul_indices: List[Tensor], intermediate_results: List[Tensor], use_double_precision: bool = True) → Tensor

Compute gradients for local skeleton state through backpropagation.

This function implements the backward pass for forward kinematics, computing gradients of the loss with respect to local joint states given gradients with respect to global joint states.

The backpropagation uses the intermediate results saved during the forward pass to efficiently compute gradients without recomputing the full forward kinematics chain.

Gradient Flow: For each joint j with parent p, the backward pass computes: ∂L/∂local_j = ∂L/∂global_j * ∂global_j/∂local_j

Where the Jacobian ∂global_j/∂local_j depends on the parent’s global state and is computed using the chain rule through similarity transformation composition.

Args:

global_skel_state: Global joint states from forward pass, shape (batch_size, num_joints, 8).: These states are modified in-place during backpropagation.
grad_global_skel_state: Gradients w.r.t. global joint states, shape (batch_size, num_joints, 8).: Input gradients from downstream computations (e.g., skinning loss).
prefix_mul_indices: List of [child_index, parent_index] tensor pairs that defined: the forward pass traversal order. Backprop processes these in reverse.
intermediate_results: List of intermediate joint states saved during forward pass.: Required to compute accurate gradients without numerical drift.
use_double_precision: If True, performs computations in float64 for numerical stability.: Recommended for deep kinematic chains and precise gradient computation.

Returns:

grad_local_skel_state: Gradients w.r.t. local joint states, shape (batch_size, num_joints, 8).: These gradients can be used to update joint parameters via optimization.

Note:

This function processes the kinematic chain in reverse order of the forward pass, accumulating gradients from children to parents while reconstructing intermediate states.

See Also:

global_skel_state_from_local_skel_state_impl(): Forward pass implementation

pymomentum.backend.skel_state_backend.global_skel_state_from_local_skel_state_no_grad(local_skel_state: Tensor, prefix_mul_indices: List[Tensor], save_intermediate_results: bool = True, use_double_precision: bool = True) → Tuple[Tensor, List[Tensor]]

Compute global skeleton state without gradient tracking.

This is a convenience wrapper around global_skel_state_from_local_skel_state_impl that explicitly disables gradient computation using torch.no_grad(). Useful for inference-only forward passes to reduce memory usage.

Args:: local_skel_state: Local joint transformations, shape (batch_size, num_joints, 8) prefix_mul_indices: List of [child_index, parent_index] tensor pairs save_intermediate_results: Whether to save intermediate states for backprop use_double_precision: Whether to use float64 for numerical stability
Returns:: global_skel_state: Global joint transformations, shape (batch_size, num_joints, 8) intermediate_results: List of intermediate joint states from forward pass
See Also:: global_skel_state_from_local_skel_state_impl(): Implementation function

pymomentum.backend.skel_state_backend.unpose_from_momentum_global_joint_state(verts: Tensor, global_joint_state: Tensor, binded_joint_state_inv: Tensor, skin_indices_flattened: Tensor, skin_weights_flattened: Tensor, vert_indices_flattened: Tensor, with_high_precision: bool = True) → Tensor

The inverse function of skinning(). WARNING: the precision is low…

Args:: verts: [batch_size, num_verts, 3] global_joint_state (th.Tensor): (B, J, 8) binded_joint_state (th.Tensor): (J, 8) skin_indices_flattened: (N, ) LBS skinning nbr joint indices skin_weights_flattened: (N, ) LBS skinning nbr joint weights vert_indices_flattened: (N, ) LBS skinning nbr corresponding vertex indices with_high_precision: if True, use high precision solver (LDLT),

but requires a cuda device sync

pymomentum.backend.trs_backend module

TRS Backend for PyMomentum

This module provides efficient forward kinematics and skinning operations using the TRS (Translation-Rotation-Scale) representation where each transformation component is stored separately.

The TRS representation uses separate tensors for translation (3D), rotation matrices (3x3), and scale factors (1D), making it suitable for applications that need explicit access to individual transformation components.

Performance Notes: This backend is typically 25-50% faster than the skeleton state backend in PyTorch, likely due to not requiring quaternion normalization operations. While it doesn’t match the C++ reference implementation exactly (use skel_state_backend for that), it provides excellent performance for PyTorch-based applications.

Key Functions: - global_trs_state_from_local_trs_state: Forward kinematics from local to global joint states - skin_points_from_trs_state: Linear blend skinning using TRS transformations - local_trs_state_from_joint_params: Convert joint parameters to local TRS states

Related Modules: - skel_state_backend: Alternative backend using compact 8-parameter skeleton states - trs: Core TRS transformation operations and utilities

class pymomentum.backend.trs_backend.ForwardKinematicsFromLocalTransformationJIT(*args, **kwargs)

Bases: Function

static backward(ctx, grad_joint_state_t: Tensor, grad_joint_state_r: Tensor, grad_joint_state_s: Tensor, _0) → Tuple[Tensor, Tensor, Tensor, None]

Define a formula for differentiating the operation with backward mode automatic differentiation.

This function is to be overridden by all subclasses. (Defining this function is equivalent to defining the vjp function.)

It must accept a context ctx as the first argument, followed by as many outputs as the forward() returned (None will be passed in for non tensor outputs of the forward function), and it should return as many tensors, as there were inputs to forward(). Each argument is the gradient w.r.t the given output, and each returned value should be the gradient w.r.t. the corresponding input. If an input is not a Tensor or is a Tensor not requiring grads, you can just pass None as a gradient for that input.

The context can be used to retrieve tensors saved during the forward pass. It also has an attribute ctx.needs_input_grad as a tuple of booleans representing whether each input needs gradient. E.g., backward() will have ctx.needs_input_grad[0] = True if the first input to forward() needs gradient computed w.r.t. the output.

static forward(local_state_t: Tensor, local_state_r: Tensor, local_state_s: Tensor, prefix_mul_indices: List[Tensor]) → Tuple[Tensor, Tensor, Tensor, List[Tuple[Tensor, Tensor, Tensor]]]

Define the forward of the custom autograd Function.

This function is to be overridden by all subclasses. There are two ways to define forward:

Usage 1 (Combined forward and ctx):

@staticmethod
def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any:
    pass

It must accept a context ctx as the first argument, followed by any number of arguments (tensors or other types).
See combining-forward-context for more details

Usage 2 (Separate forward and ctx):

@staticmethod
def forward(*args: Any, **kwargs: Any) -> Any:
    pass

@staticmethod
def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None:
    pass

The forward no longer accepts a ctx argument.
Instead, you must also override the torch.autograd.Function.setup_context() staticmethod to handle setting up the ctx object. output is the output of the forward, inputs are a Tuple of inputs to the forward.
See extending-autograd for more details

The context can be used to store arbitrary data that can be then retrieved during the backward pass. Tensors should not be stored directly on ctx (though this is not currently enforced for backward compatibility). Instead, tensors should be saved either with ctx.save_for_backward() if they are intended to be used in backward (equivalently, vjp) or ctx.save_for_forward() if they are intended to be used for in jvp.

static setup_context(ctx, inputs, outputs) → None

There are two ways to define the forward pass of an autograd.Function.

Either:

Override forward with the signature forward(ctx, *args, **kwargs). setup_context is not overridden. Setting up the ctx for backward happens inside the forward.
Override forward with the signature forward(*args, **kwargs) and override setup_context. Setting up the ctx for backward happens inside setup_context (as opposed to inside the forward)

See torch.autograd.Function.forward() and extending-autograd for more details.

pymomentum.backend.trs_backend.global_trs_state_from_local_trs_state(local_state_t: Tensor, local_state_r: Tensor, local_state_s: Tensor, prefix_mul_indices: List[Tensor]) → Tuple[Tensor, Tensor, Tensor]

Compute global TRS state from local joint transformations (user-facing wrapper).

This is the main entry point for forward kinematics using TRS states. It automatically selects between JIT-compiled and autograd-enabled implementations based on the execution context.

Args:: local_state_t: Local joint translations, shape (batch_size, num_joints, 3). local_state_r: Local joint rotations, shape (batch_size, num_joints, 3, 3). local_state_s: Local joint scales, shape (batch_size, num_joints, 1). prefix_mul_indices: List of [child_index, parent_index] tensor pairs defining

the kinematic hierarchy traversal order.
Returns:: global_state_t: Global joint translations, shape (batch_size, num_joints, 3). global_state_r: Global joint rotations, shape (batch_size, num_joints, 3, 3). global_state_s: Global joint scales, shape (batch_size, num_joints, 1).
Note:: When called within torch.jit.script or torch.jit.trace context, uses the JIT-compiled implementation for maximum performance. Otherwise, uses the autograd-enabled version for gradient computation.
See Also:: global_trs_state_from_local_trs_state_impl(): JIT implementation local_trs_state_from_joint_params(): Convert joint parameters to local states

pymomentum.backend.trs_backend.global_trs_state_from_local_trs_state_forward_only(local_state_t: Tensor, local_state_r: Tensor, local_state_s: Tensor, prefix_mul_indices: list[Tensor]) → tuple[Tensor, Tensor, Tensor]

Compute global TRS state from local joint transformations (forward-only wrapper).

This is a forward-only version that bypasses autograd completely, used when gradients are not needed and maximum performance is required.

Args:: local_state_t: Local joint translations, shape (batch_size, num_joints, 3). local_state_r: Local joint rotations, shape (batch_size, num_joints, 3, 3). local_state_s: Local joint scales, shape (batch_size, num_joints, 1). prefix_mul_indices: List of [child_index, parent_index] tensor pairs.
Returns:: global_state_t: Global joint translations, shape (batch_size, num_joints, 3). global_state_r: Global joint rotations, shape (batch_size, num_joints, 3, 3). global_state_s: Global joint scales, shape (batch_size, num_joints, 1).
See Also:: global_trs_state_from_local_trs_state(): Main user-facing function with autograd

pymomentum.backend.trs_backend.unpose_from_global_joint_state(verts: Tensor, t: Tensor, r: Tensor, s: Tensor, t0: Tensor, r0: Tensor, skin_indices_flattened: Tensor, skin_weights_flattened: Tensor, vert_indices_flattened: Tensor, with_high_precision: bool = True) → Tensor

The inverse function of skinning(). WARNING: the precision is low…

Args:: verts: [batch_size, num_verts, 3] t: (B, J, 3) Translation of the joints r: (B, J, 3, 3) Rotation of the joints s: (B, J, 1) Scale of the joints t0: (J, 3) Translation of inverse bind pose r0: (J, 3, 3) Rotation of inverse bind pose skin_indices_flattened: (N, ) LBS skinning nbr joint indices skin_weights_flattened: (N, ) LBS skinning nbr joint weights vert_indices_flattened: (N, ) LBS skinning nbr corresponding vertex indices with_high_precision: if True, use high precision solver (LDLT),

but requires a cuda device sync

pymomentum.backend.utils module

Backend Utility Functions for PyMomentum

This module provides utility functions that are specific to backend operations and were previously available in real_lbs_pytorch but are now implemented within pymomentum for the backend porting effort.

class pymomentum.backend.utils.LBSAdapter(character: Character, device: str = 'cpu')

Bases: object

Adapter class to make pymomentum Character compatible with LBS interface.

This adapter provides a compatibility layer that allows pymomentum Character objects to be used in tests and code that originally expected the LBS interface from the real_lbs_pytorch library. It extracts and converts the necessary properties from a Character object into the expected tensor format and provides the methods that the original LBS interface had.

__init__(character: Character, device: str = 'cpu') → None

Initialize the LBS adapter with a pymomentum Character.

Parameters:

character – The pymomentum Character to adapt
device – Device to place tensors on (default: “cpu”)

assemble_pose_and_scale_(pose: Tensor, scale: Tensor) → Tensor

Assemble pose and scale parameters into model parameters.

Parameters:

pose – Pose parameters tensor
scale – Scale parameters tensor

Returns:

Combined model parameters tensor

property device: device: Get the device that tensors are stored on.

to(device: device) → LBSAdapter

Move adapter to specified device.

Parameters:: device – Target device to move tensors to
Returns:: New LBSAdapter instance on the target device

pymomentum.backend.utils.calc_fk_prefix_multiplication_indices(joint_parents: Tensor) → List[Tensor]

Calculate prefix multiplication indices for forward kinematics.

This function computes the indices needed for efficient prefix multiplication during forward kinematics computation. The algorithm builds kinematic chains for each joint and determines the multiplication order for parallel processing.

Parameters:: joint_parents (torch.Tensor) – Parent joint index for each joint. For root joint, its parent is -1.
Returns:: List of prefix multiplication indices per level. For each level, index[0] is the source and index[1] is the target indices.
Return type:: List[torch.Tensor]

pymomentum.backend.utils.flatten_skinning_weights_and_indices(skin_weights: Tensor, skin_indices: Tensor) → Tuple[Tensor, Tensor, Tensor]

Decompress LBS skinning weights and indices into flattened arrays.

This function takes the typical (V, 8) sparse representation of skinning weights and indices and converts them into flattened arrays by removing zero weights, making them suitable for efficient skinning computation.

Parameters:

skin_weights (torch.Tensor) – Skinning weights tensor of shape (V, 8).
skin_indices (torch.Tensor) – Skinning joint indices tensor of shape (V, 8).

Returns:

Tuple of (skin_indices_flattened, skin_weights_flattened, vert_indices_flattened).

Return type:

Tuple[torch.Tensor, torch.Tensor, torch.Tensor]