Analogous to torch.unsqueeze. The dimension is relative to the shared dimensions of the rotation/translation.

Analogous to the to() method of torch Tensors Args: device: A torch device

Concatenates transformations along a new dimension. Args: ts: A list of T objects

Returns a copy of the Rotation whose underlying Tensor has been detached from its torch graph. Returns:

Updates the stored parameters using the state dict of the provided module. The module should have the same structure as that used to

Applies the transformation to a coordinate tensor. Args: pts: A [*, 3] coordinate tensor. Return

Returns the underlying rotation as a rotation matrix tensor. Returns: The rotation as a rotation matrix tens

Getter for the translation. Returns: The stored translation

Updates the stored parameters using the state dict of the provided module. The module should have the same structure as that used to

Composes the current rigid object with another. Args: r: Another Rigid object

Getter for the rotation. Returns: The rotation object

Converts a transformation to a tensor with 7 final columns, four for the quaternion followed by three for the translation.

Args: in_dim: The final dimension of inputs to the layer out_dim: The final dimension

Returns the underlying rotation as a quaternion tensor. Depending on whether the Rotation was initialized with a

Negative log probabilities

Implements the "chunking" procedure described in section 1.11.8. Layer outputs and inputs are assumed to be simple "pytrees," consisting

Constructs a transformation from a homogenous transformation tensor. Args: t: [*, 4, 4] homogeno

Right pads dimension of numpy array. Args: x: numpy like array to pad. max_len: desired length after padding pad_idx: dim

Args: c_in: Input channel dimension c_hidden: Hidden channel dimension no

Implements algorithm 21. Constructs transformations from sets of 3 points using the Gram-Schmidt algorithm. Arg

Inverts the transformation. Returns: The inverse transformation.

Generates rotation vector(s) from IGSO(3). Args: t: continuous time in [0, 1]. n_sample: number of samples to generat

Args: c_m: MSA channel dimension n: Factor multiplied to c_m to obtain the hidden cha

Args: c_z: Input channel dimension no_bins: Number of bins

Args: c_q: Input dimension of query data c_k: Input dimension of key data

input: x = PDB filename atoms = atoms to extract (optional) output: (length, atoms, coords=(x,y,z)), sequence

Queries the database using HHblits.

Convert rotation vectors to rotation matrix representation. The length of the rotation vector is the angle of rotation, the unit vector the r

Calculates the index for discretized t during IGSO(3) initialization.

Converts a transformation to a homogenous transformation tensor. Returns: A [*, 4, 4] homogenous transformat

Args: tf_dim: Final dimension of the target features msa_dim: Final dimension of the

Create a rigid 4x4 transformation matrix from two axes and transl.

Create a rigid 4x4 transformation matrix from two axes and transl.

Returns the device on which the Rigid's tensors are located. Returns: The device on which the Rigid's tensor

Truncated sum of IGSO(3) distribution. This function approximates the power series in equation 5 of "DENOISING DIFFUSION PROBABILISTIC MODELS

v: [..., 3] hat_v: [..., 3, 3]

Extracts loop associated with a prefix from mmCIF data as a list. Reference for loop_ in mmCIF: http://mmcif.wwpdb.org/docs/tutorials/mecha

Extracts loop associated with a prefix from mmCIF data as a list. Reference for loop_ in mmCIF: http://mmcif.wwpdb.org/docs/tutorials/mechanics

Matrix multiply two rotation matrices with leading dimensions.

Take the transpose of the last two dimensions.

Extract \sigma(t) corresponding to chosen sigma schedule.

Helper function to go from time t to corresponding sigma_idx.

Args: in_dim: The final dimension of inputs to the layer out_dim: The final dimension

Args: c_in: Input channel dimension c_hidden: Per-head hidden channel dimension

Args: c_t: Template embedding channel dimension c_hidden_tri_att: Per-head hidden dim

Generate a 3 by 3 identity matrix and broadcast it to a batch of target matrices. Args: target (torch.Tensor): Batch of target 3 by

Produces an ordered sequence of tensor slices that, when used in sequence on a tensor with shape dims, yields tensors that contain e

Applies a Tensor -> Tensor function to the stored translation. Args: fn: A function of

Computes FAPE loss. Args: pred_frames: [*, N_frames] Rigid object of predicted frames target