lloca.framesnet.frames.LowerIndicesFrames

class lloca.framesnet.frames.LowerIndicesFrames(frames)[source]

Bases: Frames

Frames with lower indices, obtained by multiplying with the metric. Used in LLoCaAttention to lower the key indices.

There are 2 ways to initialize an Frames object, with different arguments: - From matrices: Set matrices and optionally is_global, det, inv - As identity: Set is_identity=True, shape, device, dtype

Parameters:

  • matrices (torch.tensor) – Transformation matrices of shape (…, 4, 4)

  • is_global (bool) – Whether frames are the same for all particles in the point cloud

  • inv (torch.Tensor) – Optional cached inverse of shape (…, 4, 4). If not given, takes a bit of extra time to compute.

  • det (torch.Tensor) – Optional cached determinant of shape (…). If not given, takes a bit of extra time to compute.

  • parity (torch.Tensor) – Optional sign of shape (…) that parity-odd representations pick up, see the parity attribute. If not given, defaults to sign(det).

  • is_identity (bool) – Sets matrices to diagonal. Takes precedence over matrices: if both are given, matrices is only used to infer shape, device and dtype.

  • shape (List[int]) – Specifies matrices.shape[:-2] if is_identity. Otherwise inferred from matrices.

  • device (str) – Specifies device if is_identity. Otherwise inferred from matrices.

  • dtype (torch.dtype) – Specifies dtype if is_identity. Otherwise inferred from matrices.