tensorclouds.nn

Submodules

• tensorclouds.nn.attention
• tensorclouds.nn.egnn
• tensorclouds.nn.embed
• tensorclouds.nn.feed_forward
• tensorclouds.nn.layer_norm
• tensorclouds.nn.pooling
• tensorclouds.nn.radius_graph
• tensorclouds.nn.residual
• tensorclouds.nn.self_interaction
• tensorclouds.nn.sequence
• tensorclouds.nn.spatial
• tensorclouds.nn.tensor_fields
• tensorclouds.nn.transformer
• tensorclouds.nn.utils

Classes

TensorCloud

A TensorCloud is a collection of tensors with associated coordinates and masks.

Package Contents

class tensorclouds.nn.TensorCloud

A TensorCloud is a collection of tensors with associated coordinates and masks. It is used to represent a cloud of tensors in a 3D space, where each tensor can have different irreducible representations (irreps).

irreps_array

The array of tensors with their irreducible representations.

Type:

e3nn.IrrepsArray

mask_irreps_array

A boolean mask indicating which irreps are present.

Type:

jax.Array

coord

The coordinates of the tensors in 3D space.

Type:

jax.Array

mask_coord

A boolean mask indicating which coordinates are valid.

Type:

jax.Array

label

An optional label for the TensorCloud, e.g., for classification tasks.

Type:

jax.Array, optional

irreps_array: e3nn_jax.IrrepsArray

mask_irreps_array: jax.Array

coord: jax.Array

mask_coord: jax.Array

label: jax.Array = None

property shape

classmethod empty(irreps: e3nn_jax.Irreps)

classmethod ones_like(tc)

classmethod zeros_like(tc)

property irreps: e3nn_jax.Irreps

property mask: jax.Array

replace(**kwargs)

classmethod concatenate(clouds: List[TensorCloud]) → TensorCloud

classmethod zeros(irreps: e3nn_jax.Irreps, shape: tuple) → TensorCloud

classmethod stack(clouds: List[TensorCloud]) → TensorCloud

split(piece_sizes: List[int]) → List[TensorCloud]

centralize()

plot(view=None, viewer=None, colors=None, radius=0.04, mid=0.95)

dot(other)

norm()

Compute the norm of the irreps_array and coord of the TensorCloud. :returns: A tuple containing the norms of the irreps_array and coord. :rtype: Tuple[jax.Array, jax.Array]