63 lines
1.6 KiB
Python
63 lines
1.6 KiB
Python
from typing import Any, Dict
|
|
import torch
|
|
|
|
|
|
class PtMaskedSelect:
|
|
"""
|
|
Extracts elements from the input tensor whose corresponding value in `masked_tens` is `True`.
|
|
|
|
Example:
|
|
```
|
|
Input Tensor:
|
|
[[1, 2, 3],
|
|
[40, 50, 60],
|
|
[700, 800, 900]]
|
|
|
|
Masked Tensor:
|
|
[[True, False, False],
|
|
[False, True, False],
|
|
[False, False, True]]
|
|
|
|
Output Tensor:
|
|
[1, 50, 900]
|
|
```
|
|
|
|
category: PyTorch wrapper - Indexing and Slicing Operations
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> Dict[str, Any]:
|
|
"""
|
|
Defines the input types for the function.
|
|
|
|
Returns:
|
|
Dict[str, Any]: A dictionary specifying required input types.
|
|
"""
|
|
return {
|
|
"required": {
|
|
"tens": ("TENSOR", {}),
|
|
"masked_tens": ("TENSOR", {})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES: tuple = ("TENSOR",)
|
|
FUNCTION: str = "f"
|
|
CATEGORY: str = "Data Analysis"
|
|
|
|
def f(self, tens: torch.Tensor, masked_tens: torch.Tensor) -> tuple:
|
|
"""
|
|
Extracts elements from `tens` where `masked_tens` is `True`.
|
|
|
|
Args:
|
|
tens (torch.Tensor): The input tensor.
|
|
masked_tens (torch.Tensor): A boolean tensor specifying which elements to select.
|
|
|
|
Returns:
|
|
tuple: A tuple containing the selected tensor.
|
|
"""
|
|
if masked_tens.dtype != torch.bool:
|
|
raise ValueError("masked_tens must be a boolean tensor (dtype=torch.bool)")
|
|
|
|
result = torch.masked_select(tens, masked_tens)
|
|
return (result,)
|