62 lines
2.2 KiB
Python
62 lines
2.2 KiB
Python
from typing import Any, Dict
|
|
import torch
|
|
|
|
class PtWhere:
|
|
"""
|
|
Generates a new tensor by selecting values based on a condition tensor.
|
|
|
|
* `condition_tens` (torch.Tensor): Boolean tensor that determines which values are selected.
|
|
* `true_tens` (torch.Tensor): Values to use where the condition is `True`.
|
|
* `false_tens` (torch.Tensor): Values to use where the condition is `False`.
|
|
|
|
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": {
|
|
"condition_tens": ("TENSOR", {}),
|
|
"true_tens": ("TENSOR", {}),
|
|
"false_tens": ("TENSOR", {}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES: tuple = ("TENSOR",)
|
|
FUNCTION: str = "f"
|
|
CATEGORY: str = "Data Analysis"
|
|
|
|
def f(self,
|
|
condition_tens: torch.Tensor,
|
|
true_tens: torch.Tensor,
|
|
false_tens: torch.Tensor) -> tuple:
|
|
"""
|
|
Selects values from `true_tens` or `false_tens` based on `condition_tens`.
|
|
|
|
Args:
|
|
condition_tens (torch.Tensor): Boolean tensor indicating selection conditions.
|
|
true_tens (torch.Tensor): Values used where condition is True.
|
|
false_tens (torch.Tensor): Values used where condition is False.
|
|
|
|
Returns:
|
|
tuple: A tuple containing the modified tensor.
|
|
"""
|
|
# Ensure condition_tens is a boolean tensor
|
|
if condition_tens.dtype != torch.bool:
|
|
raise TypeError(f"Expected `condition_tens` to be a boolean tensor, got {condition_tens.dtype}")
|
|
|
|
# Ensure all tensors are compatible in shape
|
|
if not torch.broadcast_shapes(condition_tens.shape, true_tens.shape, false_tens.shape):
|
|
raise ValueError(f"Shape mismatch: condition {condition_tens.shape}, true {true_tens.shape}, false {false_tens.shape}")
|
|
|
|
# Perform the conditional selection
|
|
tens_out = torch.where(condition_tens, true_tens, false_tens)
|
|
|
|
return (tens_out,)
|