70 lines
2.5 KiB
Python
70 lines
2.5 KiB
Python
from typing import Any, Dict
|
|
import torch
|
|
import ast
|
|
|
|
class PtScatter:
|
|
"""
|
|
Generates a new tensor by replacing values at specified positions using an index tensor.
|
|
|
|
* Target tensor (tens): Some of whose values are replaced.
|
|
* Source tensor (src): Values from this tensor replace a part or whole of the target tensor.
|
|
* Index tensor (index): Specifies the positions in the target tensor for replacement.
|
|
|
|
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", {}),
|
|
"dim": ("INT", {"default": 0, "min": -10, "max": 10}),
|
|
"index": ("STRING", {"multiline": True}),
|
|
"src": ("STRING", {"multiline": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES: tuple = ("TENSOR",)
|
|
FUNCTION: str = "f"
|
|
CATEGORY: str = "Data Analysis"
|
|
|
|
def f(self,
|
|
tens: torch.Tensor,
|
|
dim: int,
|
|
index: str,
|
|
src: str) -> tuple:
|
|
"""
|
|
Replaces elements in the target tensor along the specified dimension using an index tensor.
|
|
|
|
Args:
|
|
tens (torch.Tensor): The target tensor whose values will be replaced.
|
|
dim (int): The dimension along which values are replaced.
|
|
index (str): A string representation of the index tensor (list format).
|
|
src (str): A string representation of the source tensor (list format).
|
|
|
|
Returns:
|
|
tuple: A tuple containing the modified tensor.
|
|
"""
|
|
# Parse the input index and source tensors from string representations
|
|
index_list = ast.literal_eval(index)
|
|
src_list = ast.literal_eval(src)
|
|
|
|
# Convert to tensors
|
|
index_tensor = torch.tensor(index_list, dtype=torch.long, device=tens.device)
|
|
src_tensor = torch.tensor(src_list, dtype=tens.dtype, device=tens.device)
|
|
|
|
# Ensure index and src shapes match
|
|
if index_tensor.shape != src_tensor.shape:
|
|
raise ValueError(f"Shape mismatch: index tensor {index_tensor.shape} and src tensor {src_tensor.shape} must match.")
|
|
|
|
# Perform the scatter operation
|
|
updated_tens = torch.scatter(tens.clone(), dim=dim, index=index_tensor, src=src_tensor)
|
|
|
|
return (updated_tens,)
|