nodes/HowToSD-ComfyUI-Data-Analysis
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Added more PyTorch wrapper nodes.

Browse Source
This commit is contained in:
HowToSD
2025-02-27 23:52:21 -08:00
parent b0782fb422
commit b22193dc2b
58 changed files with 2753 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/pytorch_wrapper/pt_argmax.py
+59
View File
@@ -0,0 +1,59 @@
from typing import Any, Dict
import torch
from .utils import str_to_dim
class PtArgmax:
"""
Computes the indices of the maximum values of a PyTorch tensor along the specified dimension(s).
Specify the dimension in the `dim` field using an integer as shown below:
```
0
```
category: PyTorch wrapper - Reduction operation & Summary statistics
"""
@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": ("STRING", {"multiline": True}),
"keepdim": ("BOOLEAN", {"default": False})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, dim: str, keepdim: bool) -> tuple:
"""
Computes the indices of the maximum values along the specified dimension(s).
Args:
tens (torch.Tensor): The input tensor.
dim (str): A string representation of the dimension(s) along which to compute argmax,
e.g., "0" or "".
keepdim (bool): Whether to retain reduced dimensions.
Returns:
tuple: A tuple containing the tensor with indices of maximum values.
"""
dim = str_to_dim(dim)
# Ensure keepdim is a boolean
if not isinstance(keepdim, bool):
raise TypeError(f"Invalid type for keepdim: {keepdim}. Must be a boolean.")
# Compute the argmax along the specified dimension(s)
tens = torch.argmax(tens, dim=dim, keepdim=keepdim)
return (tens,)
modules/pytorch_wrapper/pt_argmin.py
+59
View File
@@ -0,0 +1,59 @@
from typing import Any, Dict
import torch
from .utils import str_to_dim
class PtArgmin:
"""
Computes the indices of the minimum values of a PyTorch tensor along the specified dimension(s).
Specify the dimension in the `dim` field using an integer as shown below:
```
0
```
category: PyTorch wrapper - Reduction operation & Summary statistics
"""
@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": ("STRING", {"multiline": True}),
"keepdim": ("BOOLEAN", {"default": False})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, dim: str, keepdim: bool) -> tuple:
"""
Computes the indices of the minimum values along the specified dimension(s).
Args:
tens (torch.Tensor): The input tensor.
dim (str): A string representation of the dimension(s) along which to compute argmin,
e.g., "0" or "".
keepdim (bool): Whether to retain reduced dimensions.
Returns:
tuple: A tuple containing the tensor with indices of minimum values.
"""
dim = str_to_dim(dim)
# Ensure keepdim is a boolean
if not isinstance(keepdim, bool):
raise TypeError(f"Invalid type for keepdim: {keepdim}. Must be a boolean.")
# Compute the argmin along the specified dimension(s)
tens = torch.argmin(tens, dim=dim, keepdim=keepdim)
return (tens,)
modules/pytorch_wrapper/pt_gather.py
+94
View File
@@ -0,0 +1,94 @@
from typing import Any, Dict
import torch
import ast
class PtGather:
"""
Generates a tensor based on the index tensor using PyTorch's `gather` function.
Specify the `dim` field using an integer to indicate the axis along which to gather elements.
Each input tensor element is selected based on the index tensor.
For example, consider the following input tensor:
```
[ [10, 20, 30, 40, 50],
[100, 200, 300, 400, 500]]
```
and the corresponding index tensor:
```
[ [0, 4, 3, 2, 0],
[4, 3, 0, 0, 0]]
```
If `dim = 0`, the function scans row-wise, treating each index value as a row index.
If `dim = 1`, it scans column-wise, interpreting each index value as a column index.
Let's analyze the case where `dim = 1`:
* Start at row 0, column 0.
* The index at this position is `index[0, 0] = 0`.
* PyTorch selects the value at position 0 from the array `[10, 20, 30, 40, 50]`, which is `10`.
* The result at this stage is:
```
[[10, ?, ?, ?, ?]
[?, ?, ?, ?, ?]]
```
* Moving to column 1: `index[0, 1] = 4`.
* The value at index 4 in `[10, 20, 30, 40, 50]` is `50`.
* The updated result is:
```
[[10, 50, ?, ?, ?]
[?, ?, ?, ?, ?]]
```
* Repeating this process for all elements constructs the final gathered tensor.
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})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, dim: int, index: str) -> tuple:
"""
Gathers elements along the specified dimension using an index tensor.
Args:
tens (torch.Tensor): The input tensor.
dim (int): The dimension along which to gather values.
index (str): A string representation of the index tensor, provided in list format.
Returns:
tuple: A tuple containing the gathered tensor.
"""
# Parse the index string to a Python list
index = ast.literal_eval(index)
# Convert the parsed list to a tensor
index_tensor = torch.tensor(index, dtype=torch.long, device=tens.device)
# Perform the gather operation
gathered_tens = torch.gather(tens, dim=dim, index=index_tensor)
return (gathered_tens,)
modules/pytorch_wrapper/pt_index_select.py
+78
View File
@@ -0,0 +1,78 @@
from typing import Any, Dict
import torch
import ast
class PtIndexSelect:
"""
Extracts elements from the input tensor along a specified dimension using an index tensor.
`dim` specifies the dimension along which the selection occurs.
* For example, for 2D tensors:
** If `dim=0`, `index` selects rows.
** If `dim=1`, `index` selects columns (for 2D tensors).
`index` is a list of indices that will be converted into a tensor.
Example:
```
Input Tensor:
[[1, 2, 3],
[40, 50, 60],
[700, 800, 900]]
index = "[0, 2]"
dim = 0
Output Tensor:
[[1, 2, 3],
[700, 800, 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", {}),
"dim": ("INT", {"default": 0, "min": -10, "max": 10}),
"index": ("STRING", {"multiline": True})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, dim: int, index: str) -> tuple:
"""
Extracts elements from the input tensor along the specified dimension.
Args:
tens (torch.Tensor): The input tensor.
dim (int): The dimension along which elements are selected.
index (str): A string representation of a list of indices, e.g., "[0, 2]".
Returns:
tuple: A tuple containing the selected tensor.
"""
# Convert the string index to a Python list
index_list = ast.literal_eval(index)
if not isinstance(index_list, list):
raise ValueError("Index must be a list of integers.")
# Convert the index list to a tensor
index_tensor = torch.tensor(index_list, dtype=torch.long)
# Perform index selection
result = torch.index_select(tens, dim=dim, index=index_tensor)
return (result,)
modules/pytorch_wrapper/pt_masked_select.py
+62
View File
@@ -0,0 +1,62 @@
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,)
modules/pytorch_wrapper/pt_max.py
+59
View File
@@ -0,0 +1,59 @@
from typing import Any, Dict
import torch
import ast
from .utils import str_to_dim
class PtMax:
"""
Computes the maximum values of a PyTorch tensor along the specified dimension(s).
Specify the dimension(s) in the `dim` field using an integer, a list, or a tuple, as shown below:
```
0, [0], or (1, 2)
```
category: PyTorch wrapper - Reduction operation & Summary statistics
"""
@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": ("STRING", {"multiline": True}),
"keepdim": ("BOOLEAN", {"default": False})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, dim: str, keepdim: bool) -> tuple:
"""
Computes the maximum values along the specified dimension(s).
Args:
tens (torch.Tensor): The input tensor.
dim (str): A string representation of the dimension(s) along which to compute the maximum,
e.g., "0", "[0]", or "(1, 2)".
keepdim (bool): Whether to retain reduced dimensions.
Returns:
tuple: A tuple containing the tensor with maximum values.
"""
dim = str_to_dim(dim)
# Ensure keepdim is a boolean
if not isinstance(keepdim, bool):
raise TypeError(f"Invalid type for keepdim: {keepdim}. Must be a boolean.")
# Compute the max along the specified dimension(s)
tens = torch.amax(tens, dim=dim, keepdim=keepdim)
return (tens,)
modules/pytorch_wrapper/pt_mean.py
+61
View File
@@ -0,0 +1,61 @@
from typing import Any, Dict
import torch
from .utils import str_to_dim
class PtMean:
"""
Computes the mean of a PyTorch tensor along the specified dimension(s).
This node supports only floating-point or complex tensors as input.
Specify the dimension(s) in the `dim` field using an integer, a list, or a tuple, as shown below:
```
0, [0], or (1, 2)
```
category: PyTorch wrapper - Reduction operation & Summary statistics
"""
@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": ("STRING", {"multiline": True}),
"keepdim": ("BOOLEAN", {"default": False})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, dim: str, keepdim: bool) -> tuple:
"""
Computes the mean along the specified dimension(s).
Args:
tens (torch.Tensor): The input tensor.
dim (str): A string representation of the dimension(s) along which to compute the mean,
e.g., "0", "[0]", or "(1, 2)".
keepdim (bool): Whether to retain reduced dimensions.
Returns:
tuple: A tuple containing the mean tensor.
"""
dim = str_to_dim(dim)
# Ensure keepdim is a boolean
if not isinstance(keepdim, bool):
raise TypeError(f"Invalid type for keepdim: {keepdim}. Must be a boolean.")
# Compute the mean along the specified dimension(s)
tens = torch.mean(tens, dim=dim, keepdim=keepdim)
return (tens,)
modules/pytorch_wrapper/pt_median.py
+66
View File
@@ -0,0 +1,66 @@
from typing import Any, Dict
import torch
from .utils import str_to_dim
class PtMedian:
"""
Computes the median of a PyTorch tensor along the specified dimension(s).
Specify the dimension(s) in the `dim` field using an integer as shown below:
```
0
```
or
```
1
```
Note that PtMedian calls torch.median(), which returns the lower of the two numbers when the true median falls between them.
category: PyTorch wrapper - Reduction operation & Summary statistics
"""
@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": ("STRING", {"multiline": True}),
"keepdim": ("BOOLEAN", {"default": False})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, dim: str, keepdim: bool) -> tuple:
"""
Computes the median along the specified dimension(s).
Args:
tens (torch.Tensor): The input tensor.
dim (str): A string representation of the dimension(s) along which to compute the median,
e.g., "0", or "1".
keepdim (bool): Whether to retain reduced dimensions.
Returns:
tuple: A tuple containing the median tensor.
"""
dim = str_to_dim(dim)
if not isinstance(dim, int):
raise TypeError(f"Invalid type for dim: {dim}. Must be a integer.")
# Ensure keepdim is a boolean
if not isinstance(keepdim, bool):
raise TypeError(f"Invalid type for keepdim: {keepdim}. Must be a boolean.")
# Compute the median along the specified dimension(s)
tens, _ = torch.median(tens, dim=dim, keepdim=keepdim)
return (tens,)
modules/pytorch_wrapper/pt_min.py
+59
View File
@@ -0,0 +1,59 @@
from typing import Any, Dict
import torch
from .utils import str_to_dim
class PtMin:
"""
Computes the minimum values of a PyTorch tensor along the specified dimension(s).
Specify the dimension(s) in the `dim` field using an integer, a list, or a tuple, as shown below:
```
0, [0], or (1, 2)
```
category: PyTorch wrapper - Reduction operation & Summary statistics
"""
@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": ("STRING", {"multiline": True}),
"keepdim": ("BOOLEAN", {"default": False})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, dim: str, keepdim: bool) -> tuple:
"""
Computes the minimum values along the specified dimension(s).
Args:
tens (torch.Tensor): The input tensor.
dim (str): A string representation of the dimension(s) along which to compute the minimum,
e.g., "0", "[0]", or "(1, 2)".
keepdim (bool): Whether to retain reduced dimensions.
Returns:
tuple: A tuple containing the tensor with minimum values.
"""
dim = str_to_dim(dim)
# Ensure keepdim is a boolean
if not isinstance(keepdim, bool):
raise TypeError(f"Invalid type for keepdim: {keepdim}. Must be a boolean.")
# Compute the min along the specified dimension(s)
tens = torch.amin(tens, dim=dim, keepdim=keepdim)
return (tens,)
modules/pytorch_wrapper/pt_prod.py
+65
View File
@@ -0,0 +1,65 @@
from typing import Any, Dict
import torch
from .utils import str_to_dim
class PtProd:
"""
Computes the product of a PyTorch tensor along the specified dimension(s).
Specify the dimension(s) in the `dim` field using an integer as shown below:
```
0
```
or
```
1
```
category: PyTorch wrapper - Reduction operation & Summary statistics
"""
@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": ("STRING", {"multiline": True}),
"keepdim": ("BOOLEAN", {"default": False})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, dim: str, keepdim: bool) -> tuple:
"""
Computes the product along the specified dimension(s).
Args:
tens (torch.Tensor): The input tensor.
dim (str): A string representation of the dimension(s) along which to compute the product,
e.g., "0" or "1".
keepdim (bool): Whether to retain reduced dimensions.
Returns:
tuple: A tuple containing the product tensor.
"""
dim = str_to_dim(dim)
if not isinstance(dim, int):
raise TypeError(f"Invalid type for dim: {dim}. Must be a integer.")
# Ensure keepdim is a boolean
if not isinstance(keepdim, bool):
raise TypeError(f"Invalid type for keepdim: {keepdim}. Must be a boolean.")
# Compute the product along the specified dimension(s)
tens = torch.prod(tens, dim=dim, keepdim=keepdim)
return (tens,)
modules/pytorch_wrapper/pt_reshape.py
+62
View File
@@ -0,0 +1,62 @@
from typing import Any, Dict
import torch
import ast
class PtReshape:
"""
Reshapes a PyTorch tensor into a specified shape using `torch.reshape()`.
The total number of elements must remain unchanged.
For example, if a tensor has shape (2, 3, 4), and `new_shape="[6, 4]"`,
the output will have shape (6, 4).
Use -1 to automatically infer a dimension, e.g.:
```
[2, -1]
```
category: PyTorch wrapper - Transform
"""
@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", {}),
"new_shape": ("STRING", {"multiline": True})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, new_shape: str) -> tuple:
"""
Reshapes the input tensor based on the specified shape.
Args:
tens (torch.Tensor): The input tensor.
new_shape (str): A string representation of the desired shape,
e.g., "[2, -1]" or "(6, 4)".
Returns:
tuple: A tuple containing the reshaped PyTorch tensor.
"""
# Convert string to a tuple of integers
shape_order = tuple(ast.literal_eval(new_shape))
# Ensure the parsed shape is valid
if not isinstance(shape_order, tuple) or not all(isinstance(i, int) for i in shape_order):
raise ValueError(f"Invalid format for new_shape: {new_shape}. Must be a tuple or list of integers.")
# Apply reshape operation
tens = tens.reshape(*shape_order)
return (tens,)
modules/pytorch_wrapper/pt_scatter.py
+69
View File
@@ -0,0 +1,69 @@
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,)
modules/pytorch_wrapper/pt_sqrt.py
+39
View File
@@ -0,0 +1,39 @@
from typing import Any, Dict
import torch
class PtSqrt:
"""
Computes the square root of each element in a PyTorch tensor.
category: PyTorch wrapper - Math operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor) -> tuple:
"""
Computes the square root of a PyTorch tensor element-wise.
Args:
tens_a (torch.Tensor): A PyTorch tensor of numerical type.
Returns:
tuple: A tuple containing the resultant tensor after applying the square root function.
"""
return (torch.sqrt(tens_a),)
modules/pytorch_wrapper/pt_std.py
+68
View File
@@ -0,0 +1,68 @@
from typing import Any, Dict
import torch
from .utils import str_to_dim
class PtStd:
"""
Computes the standard deviation of a PyTorch tensor along the specified dimension(s).
Specify 1 to use Bessel's correction (N-1) for sample standard deviation.
Specify 0 to compute population standard deviation.
Specify the dimension(s) in the `dim` field using an integer, a list, or a tuple, as shown below:
```
0, [0], or (1, 2)
```
category: PyTorch wrapper - Reduction operation & Summary statistics
"""
@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", {}),
"correction": ("INT", {"default": 1, "min": 0, "max": 1}),
"dim": ("STRING", {"multiline": True}),
"keepdim": ("BOOLEAN", {"default": False})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, dim: str, correction: int, keepdim: bool) -> tuple:
"""
Computes the standard deviation along the specified dimension(s).
Args:
tens (torch.Tensor): The input tensor.
dim (str): A string representation of the dimension(s) along which to compute the standard deviation,
e.g., "0", "[0]", or "(1, 2)".
correction (int): 0 for population standard deviation, 1 for sample standard deviation (Bessel's correction).
keepdim (bool): Whether to retain reduced dimensions.
Returns:
tuple: A tuple containing the standard deviation tensor.
"""
dim = str_to_dim(dim)
# Ensure correction is valid (0 or 1)
if correction not in (0, 1):
raise ValueError(f"Invalid correction value: {correction}. Must be 0 (population) or 1 (sample).")
# Ensure keepdim is a boolean
if not isinstance(keepdim, bool):
raise TypeError(f"Invalid type for keepdim: {keepdim}. Must be a boolean.")
# Compute the standard deviation along the specified dimension(s)
tens = torch.std(tens, dim=dim, correction=correction, keepdim=keepdim)
return (tens,)
modules/pytorch_wrapper/pt_sum.py
+59
View File
@@ -0,0 +1,59 @@
from typing import Any, Dict
import torch
from .utils import str_to_dim
class PtSum:
"""
Computes the sum of a PyTorch tensor along the specified dimension(s).
Specify the dimension(s) in the `dim` field using an integer, a list, or a tuple, as shown below:
```
0, [0], or (1, 2)
```
category: PyTorch wrapper - Reduction operation & Summary statistics
"""
@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": ("STRING", {"multiline": True}),
"keepdim": ("BOOLEAN", {"default": False})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, dim: str, keepdim: bool) -> tuple:
"""
Computes the sum along the specified dimension(s).
Args:
tens (torch.Tensor): The input tensor.
dim (str): A string representation of the dimension(s) along which to compute the sum,
e.g., "0", "[0]", or "(1, 2)".
keepdim (bool): Whether to retain reduced dimensions.
Returns:
tuple: A tuple containing the summed tensor.
"""
dim = str_to_dim(dim)
# Ensure keepdim is a boolean
if not isinstance(keepdim, bool):
raise TypeError(f"Invalid type for keepdim: {keepdim}. Must be a boolean.")
# Compute the sum along the specified dimension(s)
tens = torch.sum(tens, dim=dim, keepdim=keepdim)
return (tens,)
modules/pytorch_wrapper/pt_var.py
+68
View File
@@ -0,0 +1,68 @@
from typing import Any, Dict
import torch
from .utils import str_to_dim
class PtVar:
"""
Computes the variance of a PyTorch tensor along the specified dimension(s).
Specify 1 to use Bessel's correction (N-1) for sample variance.
Specify 0 to compute population variance.
Specify the dimension(s) in the `dim` field using an integer, a list, or a tuple, as shown below:
```
0, [0], or (1, 2)
```
category: PyTorch wrapper - Reduction operation & Summary statistics
"""
@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", {}),
"correction": ("INT", {"default": 1, "min": 0, "max": 1}),
"dim": ("STRING", {"multiline": True}),
"keepdim": ("BOOLEAN", {"default": False})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, dim: str, correction: int, keepdim: bool) -> tuple:
"""
Computes the variance along the specified dimension(s).
Args:
tens (torch.Tensor): The input tensor.
dim (str): A string representation of the dimension(s) along which to compute the variance,
e.g., "0", "[0]", or "(1, 2)".
correction (int): 0 for population variance, 1 for sample variance (Bessel's correction).
keepdim (bool): Whether to retain reduced dimensions.
Returns:
tuple: A tuple containing the variance tensor.
"""
dim = str_to_dim(dim)
# Ensure correction is valid (0 or 1)
if correction not in (0, 1):
raise ValueError(f"Invalid correction value: {correction}. Must be 0 (population) or 1 (sample).")
# Ensure keepdim is a boolean
if not isinstance(keepdim, bool):
raise TypeError(f"Invalid type for keepdim: {keepdim}. Must be a boolean.")
# Compute the variance along the specified dimension(s)
tens = torch.var(tens, dim=dim, correction=correction, keepdim=keepdim)
return (tens,)
modules/pytorch_wrapper/pt_where.py
+61
View File
@@ -0,0 +1,61 @@
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,)
modules/pytorch_wrapper/utils.py
+39 -1
View File
@@ -1,3 +1,5 @@
import ast
from typing import Optional, Tuple, Union
import torch
DTYPE_MAPPING = {
@@ -11,4 +13,40 @@ DTYPE_MAPPING = {
"int32": torch.int32,
"int64": torch.int64,
"bool": torch.bool
}
}
def str_to_dim(dim: str) -> Optional[Union[int, Tuple[int, ...]]]:
"""
Converts a string representation of a dimension into an integer or a tuple of integers.
Args:
dim (str): A string representing an integer or a tuple of integers.
Returns:
Optional[Union[int, Tuple[int, ...]]]:
- An integer if the input represents a single integer.
- A tuple of integers if the input represents a list or tuple.
- None if the input is empty.
Raises:
ValueError: If the parsed value is not an integer or a tuple of integers.
"""
if dim:
# Convert string to an integer, list, or tuple of integers
parsed_dim = ast.literal_eval(dim)
# Convert list to tuple
if isinstance(parsed_dim, list):
parsed_dim = tuple(parsed_dim)
# Ensure the parsed dim is valid
if not isinstance(parsed_dim, (int, tuple)) or (
isinstance(parsed_dim, tuple) and not all(isinstance(i, int) for i in parsed_dim)
):
raise ValueError(f"Invalid format for dim: {parsed_dim}. Must be an integer or a tuple of integers.")
return parsed_dim
return None