TorchEmbedding¶
- class probly.representation.embedding.torch.TorchEmbedding(embeddings: Tensor = <property object>)[source]¶
Bases:
TorchAxisProtected[Tensor],Embedding[Tensor]Embedding vectors with a protected trailing embedding axis.
Shape:
batch_shape. The underlyingembeddingstensor has shape(*batch_shape, embedding_dim).- clone(*, memory_format: memory_format = torch.preserve_format) Self[source]¶
Return a copy of the array.
- cpu(memory_format: memory_format = torch.preserve_format) Self[source]¶
Move the array to the CPU.
- cuda(device: device | str | None = None, non_blocking: bool = False, memory_format: memory_format = torch.preserve_format) Self[source]¶
Move the array to the GPU.
- embeddings: torch.Tensor¶
- gather(dim: int, index: Tensor) Self[source]¶
Return a copy with gathered protected values along a batch dimension.
- permitted_functions: ClassVar[set[Callable[..., Any]]] = {<built-in method mean of type object>, <built-in method sum of type object>, <function torch_average>}¶
- permute(*dims: Size | int | tuple[int] | list[int]) Self[source]¶
Return a permuted version of the array.
- protected_values(func: Callable | None = None) dict[str, TorchProtectedValue] | None[source]¶
Return all protected field values as-is.
Optionally takes the torch function that triggered the call for context. This can be used to conditionally modify the returned values or prevent them from being accessed.
- size(dim: int | None = None) int | Size[source]¶
Return the size of the array along the given dimension.