Mode of the distribution. Args: straight_through_gradient: Boolean; if True, it uses the straight-through gradient esti

Builds the module. Args: n_caps: int, number of capsules. dim_caps: int, number of capsule parameters hidden_si

Transformer-like self-attention. Args: queries: Tensor of shape [B, N, d_k]. keys: Tensor of shape [B, M, d_k]. values: :

Computes capsule sparsity loss according to the specified type.

Computes entropy in capsule activations.

Creates a Gaussian mixture by loc(mean), scale(std) and mixing logits with K number of components. loc: tensor [B, K, ...] o

Args: capsule_layer: a capsule layer to predict object parameters

Args: normal_dist: torch normal distribution, [B, K, ...] mixing_logits: tensor [B, K, ...] with K the number of componen

Adds noise to tensors.

Computes l2 penalty on capsule activations.

Args: feature: Tensor of encodings of shape [B, O, F]. parent_transform: Tuple of (matrix, vector). parent_pres

Args: obj_encoding: Tensor of shape [B, O, D]. part_pose: Tensor of shape [B, M, P] part_presence: Tensor of sh

Args: feature: [B, n_templates, dim_feature] tensor; these features are used to change templates based on the input, if p

Builds the module. Args: templates: (B, n_templates, n_channels, *template_size) tensor pose: [B, n_templates, 6] tensor.

Multi-head transformer-like self-attention. Args: queries: Tensor of shape [B, N, d_k]. keys: Tensor of shape [B

Converts pose tensor into an affine or similarity transform. Args: pose_tensor: [..., 6] tensor. similarity (bool): nonlinear (