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

Added various PyTorch wrapper nodes to generate tensors.

Browse Source
This commit is contained in:
HowToSD
2025-02-27 13:05:17 -08:00
parent 8e36913138
commit b0782fb422
34 changed files with 1382 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+30
View File
@@ -143,6 +143,7 @@ from .modules.pytorch_wrapper.pt_floor_divide import PtFloorDiv
from .modules.pytorch_wrapper.pt_from_image import PtFromImage
from .modules.pytorch_wrapper.pt_from_latent import PtFromLatent
from .modules.pytorch_wrapper.pt_from_numpy import PtFromNumpy
from .modules.pytorch_wrapper.pt_full import PtFull
from .modules.pytorch_wrapper.pt_ge import PtGe
from .modules.pytorch_wrapper.pt_gt import PtGt
from .modules.pytorch_wrapper.pt_int_create import PtIntCreate
@@ -158,7 +159,12 @@ from .modules.pytorch_wrapper.pt_mm import PtMm
from .modules.pytorch_wrapper.pt_mul import PtMul
from .modules.pytorch_wrapper.pt_ne import PtNe
from .modules.pytorch_wrapper.pt_neg import PtNeg
from .modules.pytorch_wrapper.pt_ones import PtOnes
from .modules.pytorch_wrapper.pt_permute import PtPermute
from .modules.pytorch_wrapper.pt_pow import PtPow
from .modules.pytorch_wrapper.pt_rand import PtRand
from .modules.pytorch_wrapper.pt_rand_int import PtRandInt
from .modules.pytorch_wrapper.pt_randn import PtRandn
from .modules.pytorch_wrapper.pt_remainder import PtRemainder
from .modules.pytorch_wrapper.pt_show_text import PtShowText
from .modules.pytorch_wrapper.pt_sin import PtSin
@@ -167,6 +173,7 @@ from .modules.pytorch_wrapper.pt_size import PtSize
from .modules.pytorch_wrapper.pt_size_create import PtSizeCreate
from .modules.pytorch_wrapper.pt_size_to_numpy import PtSizeToNumpy
from .modules.pytorch_wrapper.pt_size_to_string import PtSizeToString
from .modules.pytorch_wrapper.pt_squeeze import PtSqueeze
from .modules.pytorch_wrapper.pt_sub import PtSub
from .modules.pytorch_wrapper.pt_tan import PtTan
from .modules.pytorch_wrapper.pt_tanh import PtTanh
@@ -174,6 +181,9 @@ from .modules.pytorch_wrapper.pt_to_image import PtToImage
from .modules.pytorch_wrapper.pt_to_latent import PtToLatent
from .modules.pytorch_wrapper.pt_to_numpy import PtToNumpy
from .modules.pytorch_wrapper.pt_to_rgb_tensors import PtToRgbTensors
from .modules.pytorch_wrapper.pt_unsqueeze import PtUnsqueeze
from .modules.pytorch_wrapper.pt_view import PtView
from .modules.pytorch_wrapper.pt_zeros import PtZeros
from .modules.seaborn_wrapper.sns_bar import SNSBar
from .modules.seaborn_wrapper.sns_boxplot import SNSBoxplot
from .modules.seaborn_wrapper.sns_histogram import SNSHistogram
@@ -328,6 +338,7 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
"PtFromImage": PtFromImage,
"PtFromLatent": PtFromLatent,
"PtFromNumpy": PtFromNumpy,
"PtFull": PtFull,
"PtGe": PtGe,
"PtGt": PtGt,
"PtIntCreate": PtIntCreate,
@@ -343,7 +354,12 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
"PtMul": PtMul,
"PtNe": PtNe,
"PtNeg": PtNeg,
"PtOnes": PtOnes,
"PtPermute": PtPermute,
"PtPow": PtPow,
"PtRand": PtRand,
"PtRandInt": PtRandInt,
"PtRandn": PtRandn,
"PtRemainder": PtRemainder,
"PtShowText": PtShowText,
"PtSin": PtSin,
@@ -352,6 +368,7 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
"PtSizeCreate": PtSizeCreate,
"PtSizeToNumpy": PtSizeToNumpy,
"PtSizeToString": PtSizeToString,
"PtSqueeze": PtSqueeze,
"PtSub": PtSub,
"PtTan": PtTan,
"PtTanh": PtTanh,
@@ -359,6 +376,9 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
"PtToLatent": PtToLatent,
"PtToNumpy": PtToNumpy,
"PtToRgbTensors": PtToRgbTensors,
"PtUnsqueeze": PtUnsqueeze,
"PtView": PtView,
"PtZeros": PtZeros,
"PyKvFloatCreate": PyKvFloatCreate,
"PyKvIntCreate": PyKvIntCreate,
"PyKvStringCreate": PyKvStringCreate,
@@ -516,6 +536,7 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
"PtFromImage": "Pt From Image",
"PtFromLatent": "Pt From Latent",
"PtFromNumpy": "Pt From Numpy",
"PtFull": "Pt Full",
"PtGe": "Pt Ge",
"PtGt": "Pt Gt",
"PtIntCreate": "Pt Int Create",
@@ -531,7 +552,12 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
"PtMul": "Pt Mul",
"PtNe": "Pt Ne",
"PtNeg": "Pt Neg",
"PtOnes": "Pt Ones",
"PtPermute": "Pt Permute",
"PtPow": "Pt Pow",
"PtRand": "Pt Rand",
"PtRandInt": "Pt Rand Int",
"PtRandn": "Pt Randn",
"PtRemainder": "Pt Remainder",
"PtShowText": "Pt Show Text",
"PtSin": "Pt Sin",
@@ -540,6 +566,7 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
"PtSizeCreate": "Pt Size Create",
"PtSizeToNumpy": "Pt Size To Numpy",
"PtSizeToString": "Pt Size To String",
"PtSqueeze": "Pt Squeeze",
"PtSub": "Pt Sub",
"PtTan": "Pt Tan",
"PtTanh": "Pt Tanh",
@@ -547,6 +574,9 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
"PtToLatent": "Pt To Latent",
"PtToNumpy": "Pt To Numpy",
"PtToRgbTensors": "Pt To Rgb Tensors",
"PtUnsqueeze": "Pt Unsqueeze",
"PtView": "Pt View",
"PtZeros": "Pt Zeros",
"PyKvFloatCreate": "Py Kv Float Create",
"PyKvIntCreate": "Py Kv Int Create",
"PyKvStringCreate": "Py Kv String Create",
docs/reference/node_reference.md
+13
View File
@@ -213,7 +213,13 @@
| [Pt From Image](pt_from_image.md) | Casts an Image tensor as a PyTorch tensor. |
| [Pt From Latent](pt_from_latent.md) | Casts a latent tensor as a PyTorch tensor. |
| [Pt From Numpy](pt_from_numpy.md) | Converts a NumPy ndarray to a PyTorch tensor while preserving its data type. |
| [Pt Full](pt_full.md) | Creates a PyTorch tensor filled with a specified value using the size entered in the text field. |
| [Pt Int Create](pt_int_create.md) | Creates a PyTorch tensor with 32-bit integer |
| [Pt Ones](pt_ones.md) | Creates a PyTorch tensor of ones using the size entered in the text field. |
| [Pt Rand](pt_rand.md) | Creates a PyTorch tensor with values sampled from a uniform distribution |
| [Pt Rand Int](pt_rand_int.md) | Creates a PyTorch tensor filled with random integers within a specified range using the size entered in the text field. |
| [Pt Randn](pt_randn.md) | Creates a PyTorch tensor with values sampled from a standard normal distribution (mean=0, std=1) |
| [Pt Zeros](pt_zeros.md) | Creates a PyTorch tensor of zeros using the size entered in the text field. |
## PyTorch wrapper - Tensor data conversion
| Node | Description |
| --- | --- |
@@ -221,6 +227,13 @@
| [Pt To Latent](pt_to_latent.md) | Casts a PyTorch tensor as a latent tensor. |
| [Pt To Numpy](pt_to_numpy.md) | Converts PyTorch tensor to NumPy ndarray. |
| [Pt To Rgb Tensors](pt_to_rgb_tensors.md) | Splits a PyTorch tensor into R, G, and B tensors. |
## PyTorch wrapper - Transform
| Node | Description |
| --- | --- |
| [Pt Permute](pt_permute.md) | Permutes the dimensions of a PyTorch tensor according to the specified order. |
| [Pt Squeeze](pt_squeeze.md) | Removes a dimension at the specified position in the input tensor if it is of size 1. |
| [Pt Unsqueeze](pt_unsqueeze.md) | Adds a singleton dimension at the specified position in the input tensor. |
| [Pt View](pt_view.md) | Reshapes a PyTorch tensor into a specified shape using `torch.view()`. |
## Python wrapper
| Node | Description |
| --- | --- |
docs/reference/pt_full.md
+21
View File
@@ -0,0 +1,21 @@
# Pt Full
Creates a PyTorch tensor filled with a specified value using the size entered in the text field.
For example, if you want to create a 4D tensor of batch size=2, channel=3, height=96, width=32 filled with 7, enter:
[2,3,96,32] for size and 7 for value.
## Input
| Name | Data type |
|---|---|
| value | String |
| size | String |
| data_type | |
## Output
| Data type |
|---|
| Tensor |
<HR>
Category: PyTorch wrapper - Tensor creation
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/pt_ones.md
+20
View File
@@ -0,0 +1,20 @@
# Pt Ones
Creates a PyTorch tensor of ones using the size entered in the text field.
For example, if you want to create a 4D tensor of batch size=2, channel=3, height=96, width=32, enter:
[2,3,96,32]
## Input
| Name | Data type |
|---|---|
| size | String |
| data_type | |
## Output
| Data type |
|---|
| Tensor |
<HR>
Category: PyTorch wrapper - Tensor creation
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/pt_permute.md
+25
View File
@@ -0,0 +1,25 @@
# Pt Permute
Permutes the dimensions of a PyTorch tensor according to the specified order.
For example, if a tensor has shape (2, 3, 96, 32) and `new_axes="[0, 3, 1, 2]"`,
the output will have shape (2, 32, 3, 96).
Specify the new axis order in the `new_axes` field using brackets, as shown below:
```
[0, 3, 1, 2]
```
## Input
| Name | Data type |
|---|---|
| tens | Tensor |
| new_axes | String |
## Output
| Data type |
|---|
| Tensor |
<HR>
Category: PyTorch wrapper - Transform
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/pt_rand.md
+21
View File
@@ -0,0 +1,21 @@
# Pt Rand
Creates a PyTorch tensor with values sampled from a uniform distribution
using the size entered in the text field.
For example, if you want to create a 4D tensor of batch size=2, channel=3, height=96, width=32, enter:
[2,3,96,32]
## Input
| Name | Data type |
|---|---|
| size | String |
| data_type | |
## Output
| Data type |
|---|
| Tensor |
<HR>
Category: PyTorch wrapper - Tensor creation
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/pt_rand_int.md
+22
View File
@@ -0,0 +1,22 @@
# Pt Rand Int
Creates a PyTorch tensor filled with random integers within a specified range using the size entered in the text field.
For example, if you want to create a 4D tensor of batch size=2, channel=3, height=96, width=32 with random integers in the range [0, 10),
enter [2,3,96,32] for size, 0 for min_value, and 10 for max_value.
## Input
| Name | Data type |
|---|---|
| min_value | Int |
| max_value | Int |
| size | String |
| data_type | |
## Output
| Data type |
|---|
| Tensor |
<HR>
Category: PyTorch wrapper - Tensor creation
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/pt_randn.md
+21
View File
@@ -0,0 +1,21 @@
# Pt Randn
Creates a PyTorch tensor with values sampled from a standard normal distribution (mean=0, std=1)
using the size entered in the text field.
For example, if you want to create a 4D tensor of batch size=2, channel=3, height=96, width=32, enter:
[2,3,96,32]
## Input
| Name | Data type |
|---|---|
| size | String |
| data_type | |
## Output
| Data type |
|---|
| Tensor |
<HR>
Category: PyTorch wrapper - Tensor creation
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/pt_squeeze.md
+20
View File
@@ -0,0 +1,20 @@
# Pt Squeeze
Removes a dimension at the specified position in the input tensor if it is of size 1.
Specifying a tuple of dimensions is not supported yet.
## Input
| Name | Data type |
|---|---|
| tens | Tensor |
| dim | Int |
## Output
| Data type |
|---|
| Tensor |
<HR>
Category: PyTorch wrapper - Transform
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/pt_unsqueeze.md
+19
View File
@@ -0,0 +1,19 @@
# Pt Unsqueeze
Adds a singleton dimension at the specified position in the input tensor.
For example, if a tensor has shape (3, 4) and `dim=1`, the output will have shape (3, 1, 4).
## Input
| Name | Data type |
|---|---|
| tens | Tensor |
| dim | Int |
## Output
| Data type |
|---|
| Tensor |
<HR>
Category: PyTorch wrapper - Transform
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/pt_view.md
+27
View File
@@ -0,0 +1,27 @@
# Pt View
Reshapes a PyTorch tensor into a specified shape using `torch.view()`.
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]
```
## Input
| Name | Data type |
|---|---|
| tens | Tensor |
| new_shape | String |
## Output
| Data type |
|---|
| Tensor |
<HR>
Category: PyTorch wrapper - Transform
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/pt_zeros.md
+20
View File
@@ -0,0 +1,20 @@
# Pt Zeros
Creates a PyTorch tensor of zeros using the size entered in the text field.
For example, if you want to create a 4D tensor of batch size=2, channel=3, height=96, width=32, enter:
[2,3,96,32]
## Input
| Name | Data type |
|---|---|
| size | String |
| data_type | |
## Output
| Data type |
|---|
| Tensor |
<HR>
Category: PyTorch wrapper - Tensor creation
ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
modules/pytorch_wrapper/pt_full.py
+66
View File
@@ -0,0 +1,66 @@
from typing import Any, Dict
import torch
import ast
from .utils import DTYPE_MAPPING
class PtFull:
"""
Creates a PyTorch tensor filled with a specified value using the size entered in the text field.
For example, if you want to create a 4D tensor of batch size=2, channel=3, height=96, width=32 filled with 7, enter:
[2,3,96,32] for size and 7 for value.
category: PyTorch wrapper - Tensor creation
"""
@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": {
"value": ("STRING", {"multiline": False}),
"size": ("STRING", {"multiline": True}),
"data_type": (("float32", "float16", "bfloat16", "float64", "uint8", "int8", "int16", "int32", "int64", "bool"),)
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, value: str, size: str, data_type: str) -> tuple:
"""
Creates a PyTorch tensor filled with a specified value.
Args:
value (str): String representing the value to fill the tensor.
size (str): String representation of the tensor size.
data_type (str): Desired data type of the tensor.
Returns:
tuple: A tuple containing the PyTorch tensor.
"""
# Convert string to Python list
list_size = ast.literal_eval(size)
# Convert the value to the appropriate type
dtype = DTYPE_MAPPING[data_type]
if dtype in {torch.float32, torch.float16, torch.bfloat16, torch.float64}:
fill_value = float(value)
elif dtype in {torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64}:
fill_value = int(value)
elif dtype == torch.bool:
fill_value = value.lower() in {"true", "1"}
else:
raise ValueError(f"Unsupported data type: {data_type}")
# Convert list to PyTorch Size
sz = torch.Size(list_size)
tens = torch.full(sz, fill_value, dtype=dtype)
return (tens,)
modules/pytorch_wrapper/pt_ones.py
+53
View File
@@ -0,0 +1,53 @@
from typing import Any, Dict
import torch
import ast
from .utils import DTYPE_MAPPING
class PtOnes:
"""
Creates a PyTorch tensor of ones using the size entered in the text field.
For example, if you want to create a 4D tensor of batch size=2, channel=3, height=96, width=32, enter:
[2,3,96,32]
category: PyTorch wrapper - Tensor creation
"""
@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": {
"size": ("STRING", {"multiline": True}),
"data_type": (("float32", "float16", "bfloat16", "float64", "uint8", "int8", "int16", "int32", "int64", "bool"),)
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, size: str, data_type: str) -> tuple:
"""
Creates a PyTorch tensor of ones using the size entered in the text field.
Args:
size (str): String value in the text field.
data_type (str): Desired data type of the tensor.
Returns:
tuple: A tuple containing the PyTorch tensor.
"""
# Convert string to Python list
list_size = ast.literal_eval(size)
# Convert list to PyTorch Size
sz = torch.Size(list_size)
tens = torch.ones(sz, dtype=DTYPE_MAPPING[data_type])
return (tens,)
modules/pytorch_wrapper/pt_permute.py
+60
View File
@@ -0,0 +1,60 @@
from typing import Any, Dict
import torch
import ast
class PtPermute:
"""
Permutes the dimensions of a PyTorch tensor according to the specified order.
For example, if a tensor has shape (2, 3, 96, 32) and `new_axes="[0, 3, 1, 2]"`,
the output will have shape (2, 32, 3, 96).
Specify the new axis order in the `new_axes` field using brackets, as shown below:
```
[0, 3, 1, 2]
```
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_axes": ("STRING", {"multiline": True})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, new_axes: str) -> tuple:
"""
Permutes the dimensions of the input tensor based on the specified order.
Args:
tens (torch.Tensor): The input tensor.
new_axes (str): A string representation of the desired dimension order,
e.g., "[0, 3, 1, 2]" or "(0, 3, 1, 2)".
Returns:
tuple: A tuple containing the permuted PyTorch tensor.
"""
# Convert string to a tuple of integers
permute_order = tuple(ast.literal_eval(new_axes))
# Ensure the parsed order is valid
if not isinstance(permute_order, tuple) or not all(isinstance(i, int) for i in permute_order):
raise ValueError(f"Invalid format for new_axes: {new_axes}. Must be a tuple or list of integers.")
# Apply permutation
tens = tens.permute(*permute_order)
return (tens,)
modules/pytorch_wrapper/pt_rand.py
+57
View File
@@ -0,0 +1,57 @@
from typing import Any, Dict
import torch
import ast
from .utils import DTYPE_MAPPING
class PtRand:
"""
Creates a PyTorch tensor with values sampled from a uniform distribution
using the size entered in the text field.
For example, if you want to create a 4D tensor of batch size=2, channel=3, height=96, width=32, enter:
[2,3,96,32]
category: PyTorch wrapper - Tensor creation
"""
@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": {
"size": ("STRING", {"multiline": True}),
"data_type": (("float32", "float16", "bfloat16", "float64"),)
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
@classmethod
def IS_CHANGED(cls, **kw):
return float("NaN")
def f(self, size: str, data_type: str) -> tuple:
"""
Creates a PyTorch tensor with values sampled from a uniform distribution.
Args:
size (str): String value in the text field.
data_type (str): Desired floating-point data type of the tensor.
Returns:
tuple: A tuple containing the PyTorch tensor.
"""
# Convert string to Python list
list_size = ast.literal_eval(size)
# Convert list to PyTorch Size
sz = torch.Size(list_size)
tens = torch.rand(sz, dtype=DTYPE_MAPPING[data_type])
return (tens,)
modules/pytorch_wrapper/pt_rand_int.py
+72
View File
@@ -0,0 +1,72 @@
import time
from typing import Any, Dict
import torch
import ast
import copy
from .utils import DTYPE_MAPPING
class PtRandInt:
"""
Creates a PyTorch tensor filled with random integers within a specified range using the size entered in the text field.
For example, if you want to create a 4D tensor of batch size=2, channel=3, height=96, width=32 with random integers in the range [0, 10),
enter [2,3,96,32] for size, 0 for min_value, and 10 for max_value.
category: PyTorch wrapper - Tensor creation
"""
@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": {
"min_value": ("INT", {"default": 0, "min": -2**31, "max": 2**31}),
"max_value": ("INT", {"default": 1, "min": -2**31, "max": 2**31}),
"size": ("STRING", {"multiline": True}),
"data_type": (("uint8", "int8", "int16", "int32", "int64"),)
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
@classmethod
def IS_CHANGED(cls, **kw):
return float("NaN")
def __init__(self):
# Initialize a separate generator with a different seed
self.gen = torch.Generator()
self.gen.manual_seed(int(time.time() * 1000000) % (2**32)) # Unique seed
def f(self, min_value: int, max_value: int, size: str, data_type: str) -> tuple:
"""
Creates a PyTorch tensor filled with random integers in the range [min_value, max_value).
Args:
min_value (int): Minimum value (inclusive).
max_value (int): Maximum value (exclusive).
size (str): String representation of the tensor size.
data_type (str): Desired integer data type of the tensor.
Returns:
tuple: A tuple containing the PyTorch tensor.
"""
# Convert string to Python list
list_size = ast.literal_eval(size)
# Convert list to PyTorch Size
sz = torch.Size(list_size)
# Ensure max_value is greater than min_value
if max_value <= min_value:
raise ValueError("max_value must be greater than min_value.")
# Generate random integers within the specified range
tens = torch.randint(min_value, max_value, sz, dtype=DTYPE_MAPPING[data_type], generator=self.gen)
return (tens,)
modules/pytorch_wrapper/pt_randn.py
+57
View File
@@ -0,0 +1,57 @@
from typing import Any, Dict
import torch
import ast
from .utils import DTYPE_MAPPING
class PtRandn:
"""
Creates a PyTorch tensor with values sampled from a standard normal distribution (mean=0, std=1)
using the size entered in the text field.
For example, if you want to create a 4D tensor of batch size=2, channel=3, height=96, width=32, enter:
[2,3,96,32]
category: PyTorch wrapper - Tensor creation
"""
@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": {
"size": ("STRING", {"multiline": True}),
"data_type": (("float32", "float16", "bfloat16", "float64"),)
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
@classmethod
def IS_CHANGED(cls, **kw):
return float("NaN")
def f(self, size: str, data_type: str) -> tuple:
"""
Creates a PyTorch tensor with values sampled from a standard normal distribution.
Args:
size (str): String value in the text field.
data_type (str): Desired floating-point data type of the tensor.
Returns:
tuple: A tuple containing the PyTorch tensor.
"""
# Convert string to Python list
list_size = ast.literal_eval(size)
# Convert list to PyTorch Size
sz = torch.Size(list_size)
tens = torch.randn(sz, dtype=DTYPE_MAPPING[data_type])
return (tens,)
modules/pytorch_wrapper/pt_squeeze.py
+45
View File
@@ -0,0 +1,45 @@
from typing import Any, Dict
import torch
from .utils import DTYPE_MAPPING
class PtSqueeze:
"""
Removes a dimension at the specified position in the input tensor if it is of size 1.
Specifying a tuple of dimensions is not supported yet.
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", {}),
"dim": ("INT", {"default": -1, "min": -1, "max": 10})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, dim: int) -> tuple:
"""
Removes a dimension at the specified position in the input tensor if it is of size 1.
Args:
tens (torch.Tensor): Input tensor.
dim (int): Dimension to remove if it is of size 1.
Returns:
tuple: A tuple containing the transformed PyTorch tensor.
"""
tens = torch.squeeze(tens, dim)
return (tens,)
modules/pytorch_wrapper/pt_unsqueeze.py
+44
View File
@@ -0,0 +1,44 @@
from typing import Any, Dict
import torch
class PtUnsqueeze:
"""
Adds a singleton dimension at the specified position in the input tensor.
For example, if a tensor has shape (3, 4) and `dim=1`, the output will have shape (3, 1, 4).
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", {}),
"dim": ("INT", {"default": -1, "min": -10, "max": 10})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, dim: int) -> tuple:
"""
Adds a singleton dimension at the specified position.
Args:
tens (torch.Tensor): Input tensor.
dim (int): Dimension at which to insert the new singleton dimension.
Returns:
tuple: A tuple containing the transformed PyTorch tensor.
"""
tens = torch.unsqueeze(tens, dim)
return (tens,)
modules/pytorch_wrapper/pt_view.py
+62
View File
@@ -0,0 +1,62 @@
from typing import Any, Dict
import torch
import ast
class PtView:
"""
Reshapes a PyTorch tensor into a specified shape using `torch.view()`.
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 view operation
tens = tens.view(*shape_order)
return (tens,)
modules/pytorch_wrapper/pt_zeros.py
+52
View File
@@ -0,0 +1,52 @@
from typing import Any, Dict
import torch
import ast
from .utils import DTYPE_MAPPING
class PtZeros:
"""
Creates a PyTorch tensor of zeros using the size entered in the text field.
For example, if you want to create a 4D tensor of batch size=2, channel=3, height=96, width=32, enter:
[2,3,96,32]
category: PyTorch wrapper - Tensor creation
"""
@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": {
"size": ("STRING", {"multiline": True}),
"data_type": (("float32", "float16", "bfloat16", "float64", "uint8", "int8", "int16", "int32", "int64", "bool"),)
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, size: str, data_type: str) -> tuple:
"""
Creates a PyTorch tensor of zeros using the size entered in the text field.
Args:
size (str): String value in the text field.
data_type (str): Desired data type of the tensor.
Returns:
tuple: A tuple containing the PyTorch tensor.
"""
# Convert string to Python list
list_size = ast.literal_eval(size)
# Convert list to PyTorch Size
sz = torch.Size(list_size)
tens = torch.zeros(sz, dtype=DTYPE_MAPPING[data_type])
return (tens,)
modules/pytorch_wrapper/utils.py
+14
View File
@@ -0,0 +1,14 @@
import torch
DTYPE_MAPPING = {
"float32": torch.float32,
"float16": torch.float16,
"bfloat16": torch.bfloat16,
"float64": torch.float64,
"uint8": torch.uint8,
"int8": torch.int8,
"int16": torch.int16,
"int32": torch.int32,
"int64": torch.int64,
"bool": torch.bool
}
pyproject.toml
+1 -1
View File
@@ -1,7 +1,7 @@
[project]
name = "data-analysis"
description = "Data analysis custom modules for ComfyUI - Use Pandas & Matplotlib within ComfyUI"
version = "1.4.0"
version = "1.4.1"
license = {file = "LICENSE"}
[project.urls]
tests/pytorch_wrapper/pt_full_test.py
+71
View File
@@ -0,0 +1,71 @@
import os
import sys
import unittest
import torch
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
sys.path.append(MODULE_ROOT)
from pytorch_wrapper.pt_full import PtFull
from pytorch_wrapper.utils import DTYPE_MAPPING
class TestPtFull(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.node = PtFull()
# Tensor size to test
self.size_str = "[2, 3]" # Using a small 2D tensor size for all tests
self.expected_shape = (2, 3)
# Values to test (using representative values for each type)
self.test_values = {
"float32": "1.5",
"float16": "-2.3",
"bfloat16": "3.7",
"float64": "4.9",
"uint8": "255", # Max value for uint8
"int8": "-128", # Min value for int8
"int16": "32767", # Max value for int16
"int32": "-2147483648", # Min value for int32
"int64": "9223372036854775807", # Max value for int64
"bool": "True"
}
def test_full_tensor_all_dtypes(self):
"""Test tensor creation with a specified value for all supported data types."""
for data_type, torch_dtype in DTYPE_MAPPING.items():
value_str = self.test_values[data_type]
with self.subTest(data_type=data_type, value=value_str):
result, = self.node.f(value_str, self.size_str, data_type)
# Ensure return value is a tensor
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
# Ensure result has correct shape
self.assertEqual(result.shape, torch.Size(self.expected_shape), f"Expected shape {self.expected_shape}, got {result.shape}")
# Ensure result has correct dtype
self.assertEqual(result.dtype, torch_dtype, f"Expected dtype {torch_dtype}, got {result.dtype}")
# Convert value_str to correct type
if torch_dtype in {torch.float32, torch.float16, torch.bfloat16, torch.float64}:
expected_value = float(value_str)
elif torch_dtype in {torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64}:
expected_value = int(value_str)
elif torch_dtype == torch.bool:
expected_value = value_str.lower() == "true"
else:
raise ValueError(f"Unsupported data type: {data_type}")
# Ensure all elements match the expected value
expected = torch.full(self.expected_shape, expected_value, dtype=torch_dtype)
self.assertTrue(torch.equal(result, expected), f"Tensor values do not match expected tensor for dtype {data_type}")
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/pt_ones_test.py
+44
View File
@@ -0,0 +1,44 @@
import os
import sys
import unittest
import torch
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
sys.path.append(MODULE_ROOT)
from pytorch_wrapper.pt_ones import PtOnes # Importing the PtOnes class
from pytorch_wrapper.utils import DTYPE_MAPPING # Importing the dtype mapping
class TestPtOnes(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.node = PtOnes()
# Tensor size to test
self.size_str = "[2, 3]" # Using a small 2D tensor size for all tests
self.expected_shape = (2, 3)
def test_ones_tensor_all_dtypes(self):
"""Test tensor creation with ones for all supported data types."""
for data_type, torch_dtype in DTYPE_MAPPING.items():
with self.subTest(data_type=data_type):
result, = self.node.f(self.size_str, data_type)
# Ensure return value is a tensor
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
# Ensure result has correct shape
self.assertEqual(result.shape, torch.Size(self.expected_shape), f"Expected shape {self.expected_shape}, got {result.shape}")
# Ensure result has correct dtype
self.assertEqual(result.dtype, torch_dtype, f"Expected dtype {torch_dtype}, got {result.dtype}")
# Ensure all elements are ones
expected = torch.ones(self.expected_shape, dtype=torch_dtype)
self.assertTrue(torch.equal(result, expected), f"Tensor values do not match expected ones tensor for dtype {data_type}")
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/pt_permute_test.py
+59
View File
@@ -0,0 +1,59 @@
import os
import sys
import unittest
import torch
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
sys.path.append(MODULE_ROOT)
from pytorch_wrapper.pt_permute import PtPermute
class TestPtPermute(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.node = PtPermute()
# Define test cases: (input_tensor, new_axes, expected_shape)
self.test_cases = [
# Case 1: Basic 2D permutation
(torch.ones(2, 3), "[1, 0]", (3, 2)), # (2,3) -> (3,2)
# Case 2: Permute dimensions of a 3D tensor
(torch.ones(2, 3, 4), "[2, 0, 1]", (4, 2, 3)), # (2,3,4) -> (4,2,3)
# Case 3: Identity permutation (no change)
(torch.ones(2, 3, 4), "[0, 1, 2]", (2, 3, 4)), # (2,3,4) -> (2,3,4)
# Case 4: Permute 4D tensor (like a channel-first transformation)
(torch.ones(2, 3, 96, 32), "[0, 3, 1, 2]", (2, 32, 3, 96)), # (2,3,96,32) -> (2,32,3,96)
# Case 5: Swapping first and last dimensions
(torch.ones(5, 10, 15), "[2, 1, 0]", (15, 10, 5)), # (5,10,15) -> (15,10,5)
# Case 6: Edge case with a single-dimension tensor
(torch.ones(10,), "[0]", (10,)), # (10,) -> (10,) (No change)
# Case 7: Larger 5D tensor permutation
(torch.ones(2, 3, 4, 5, 6), "[4, 3, 2, 1, 0]", (6, 5, 4, 3, 2)), # Reverse all dimensions
]
def test_permute(self):
"""Test tensor permutation operation with various dimension orders."""
for input_tensor, new_axes, expected_shape in self.test_cases:
with self.subTest(input_shape=input_tensor.shape, new_axes=new_axes):
result, = self.node.f(input_tensor, new_axes)
# Ensure return value is a tensor
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
# Ensure result has correct shape
self.assertEqual(result.shape, torch.Size(expected_shape), f"Expected shape {expected_shape}, got {result.shape}")
# Ensure result has the same dtype
self.assertEqual(result.dtype, input_tensor.dtype, "Dtype mismatch after permute operation")
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/pt_rand_int_test.py
+40
View File
@@ -0,0 +1,40 @@
import os
import sys
import unittest
import torch
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
sys.path.append(MODULE_ROOT)
from pytorch_wrapper.pt_neg import PtNeg # Updated import to PtNeg
class TestPtNeg(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.node = PtNeg()
# Test cases: (Tensor A, Expected result)
self.test_cases = [
(torch.tensor([1, -2, 3]), torch.tensor([-1, 2, -3])),
(torch.tensor([0, 5, -7]), torch.tensor([0, -5, 7])),
(torch.ones(3, 3), torch.full((3, 3), -1.0)), # Negation of ones
(torch.zeros(2, 2), torch.zeros(2, 2)), # Negation of zeros
(torch.tensor([[5], [-10]]), torch.tensor([[-5], [10]])), # Column vector case
]
def test_neg(self):
"""Test tensor element-wise negation using torch.neg()."""
for tens_a, expected in self.test_cases:
with self.subTest(tens_a=tens_a):
result, = self.node.f(tens_a)
# Ensure return value is a tensor
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
# Ensure result matches expected negation
self.assertTrue(torch.equal(result, expected), f"Expected {expected}, got {result}")
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/pt_rand_test.py
+46
View File
@@ -0,0 +1,46 @@
import os
import sys
import unittest
import torch
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
sys.path.append(MODULE_ROOT)
from pytorch_wrapper.pt_rand import PtRand # Importing the PtRand class
from pytorch_wrapper.utils import DTYPE_MAPPING # Importing the dtype mapping
class TestPtRand(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.node = PtRand()
# Tensor size to test
self.size_str = "[2, 3]" # Using a small 2D tensor size for all tests
self.expected_shape = (2, 3)
# Only floating-point types are valid for torch.rand()
self.valid_dtypes = {k: v for k, v in DTYPE_MAPPING.items() if v in {torch.float32, torch.float16, torch.bfloat16, torch.float64}}
def test_rand_tensor_all_floating_dtypes(self):
"""Test random tensor creation for all supported floating-point data types."""
for data_type, torch_dtype in self.valid_dtypes.items():
with self.subTest(data_type=data_type):
result, = self.node.f(self.size_str, data_type)
# Ensure return value is a tensor
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
# Ensure result has correct shape
self.assertEqual(result.shape, torch.Size(self.expected_shape), f"Expected shape {self.expected_shape}, got {result.shape}")
# Ensure result has correct dtype
self.assertEqual(result.dtype, torch_dtype, f"Expected dtype {torch_dtype}, got {result.dtype}")
# Ensure all elements are within range [0,1)
self.assertTrue(torch.all(result >= 0) and torch.all(result < 1), f"Tensor values are not in the range [0,1) for dtype {data_type}")
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/pt_randn_test.py
+61
View File
@@ -0,0 +1,61 @@
import os
import sys
import unittest
import torch
import ast
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
sys.path.append(MODULE_ROOT)
from pytorch_wrapper.pt_rand_int import PtRandInt # Importing the PtRandInt class
from pytorch_wrapper.utils import DTYPE_MAPPING # Importing the dtype mapping
class TestPtRandInt(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.node = PtRandInt()
# Tensor size to test
self.size_str = "[2, 3]" # Using a small 2D tensor size for all tests
self.expected_shape = (2, 3)
# Only integer types are valid for torch.randint()
self.valid_dtypes = {
k: v for k, v in DTYPE_MAPPING.items() if v in {torch.int8, torch.int16, torch.int32, torch.int64, torch.uint8}
}
# Define min and max values for each integer type
self.test_ranges = {
"int8": (-128, 127),
"int16": (-32768, 32767),
"int32": (-2147483648, 2147483647),
"int64": (-9223372036854775808, 9223372036854775807),
"uint8": (0, 255),
}
def test_randint_tensor_all_integer_dtypes(self):
"""Test random integer tensor creation for all supported integer data types."""
for data_type, torch_dtype in self.valid_dtypes.items():
min_value, max_value = self.test_ranges[data_type]
with self.subTest(data_type=data_type, min_value=min_value, max_value=max_value):
result, = self.node.f(min_value, max_value, self.size_str, data_type)
# Ensure return value is a tensor
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
# Ensure result has correct shape
self.assertEqual(result.shape, torch.Size(self.expected_shape), f"Expected shape {self.expected_shape}, got {result.shape}")
# Ensure result has correct dtype
self.assertEqual(result.dtype, torch_dtype, f"Expected dtype {torch_dtype}, got {result.dtype}")
# Ensure all elements are within range [min_value, max_value)
self.assertTrue(torch.all(result >= min_value), f"Tensor contains values less than min_value for dtype {data_type}")
self.assertTrue(torch.all(result < max_value), f"Tensor contains values greater than or equal to max_value for dtype {data_type}")
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/pt_squeeze_test.py
+57
View File
@@ -0,0 +1,57 @@
import os
import sys
import unittest
import torch
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
sys.path.append(MODULE_ROOT)
from pytorch_wrapper.pt_squeeze import PtSqueeze # Importing the PtSqueeze class
class TestPtSqueeze(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.node = PtSqueeze()
# Define test cases: (input_tensor, dim, expected_shape)
self.test_cases = [
# Case 1: Squeezing a singleton dimension
(torch.ones(1, 3, 4), 0, (3, 4)), # (1,3,4) -> (3,4)
(torch.ones(2, 1, 4), 1, (2, 4)), # (2,1,4) -> (2,4)
# Case 2: No change if dim is not a singleton
(torch.ones(2, 3, 4), 1, (2, 3, 4)), # (2,3,4) -> (2,3,4)
# Case 3: Multiple singleton dimensions, squeezing one
(torch.ones(1, 2, 1, 4), 2, (1, 2, 4)), # (1,2,1,4) -> (1,2,4)
# Case 4: Squeezing a negative index
(torch.ones(2, 3, 1), -1, (2, 3)), # (2,3,1) -> (2,3)
# Case 5: Squeezing all size 1 dimensions
(torch.ones(1, 1, 1), 0, (1, 1)), # (1,1,1) -> (1,1)
(torch.ones(1, 1, 1), 1, (1, 1)), # (1,1,1) -> (1,1)
(torch.ones(1, 1, 1), 2, (1, 1)), # (1,1,1) -> (1,1)
]
def test_squeeze(self):
"""Test tensor squeeze operation on various dimensions."""
for input_tensor, dim, expected_shape in self.test_cases:
with self.subTest(input_shape=input_tensor.shape, dim=dim):
result, = self.node.f(input_tensor, dim)
# Ensure return value is a tensor
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
# Ensure result has correct shape
self.assertEqual(result.shape, torch.Size(expected_shape), f"Expected shape {expected_shape}, got {result.shape}")
# Ensure result has the same dtype
self.assertEqual(result.dtype, input_tensor.dtype, "Dtype mismatch after squeeze operation")
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/pt_unsqueeze_test.py
+57
View File
@@ -0,0 +1,57 @@
import os
import sys
import unittest
import torch
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
sys.path.append(MODULE_ROOT)
from pytorch_wrapper.pt_unsqueeze import PtUnsqueeze # Importing the PtUnsqueeze class
class TestPtUnsqueeze(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.node = PtUnsqueeze()
# Define test cases: (input_tensor, dim, expected_shape)
self.test_cases = [
# Case 1: Adding a singleton dimension at different positions
(torch.ones(3, 4), 0, (1, 3, 4)), # (3,4) -> (1,3,4)
(torch.ones(3, 4), 1, (3, 1, 4)), # (3,4) -> (3,1,4)
(torch.ones(3, 4), 2, (3, 4, 1)), # (3,4) -> (3,4,1)
# Case 2: Adding a singleton dimension at negative indices
(torch.ones(3, 4), -1, (3, 4, 1)), # (3,4) -> (3,4,1)
(torch.ones(3, 4), -2, (3, 1, 4)), # (3,4) -> (3,1,4)
(torch.ones(3, 4), -3, (1, 3, 4)), # (3,4) -> (1,3,4)
# Case 3: Adding multiple dimensions one by one
(torch.ones(3,), 0, (1, 3)), # (3,) -> (1,3)
(torch.ones(3,), 1, (3, 1)), # (3,) -> (3,1)
(torch.ones(3,), -1, (3, 1)), # (3,) -> (3,1)
# Case 4: Expanding an already multi-dimensional tensor
(torch.ones(2, 3, 4), 1, (2, 1, 3, 4)), # (2,3,4) -> (2,1,3,4)
(torch.ones(2, 3, 4), -1, (2, 3, 4, 1)), # (2,3,4) -> (2,3,4,1)
]
def test_unsqueeze(self):
"""Test tensor unsqueeze operation on various dimensions."""
for input_tensor, dim, expected_shape in self.test_cases:
with self.subTest(input_shape=input_tensor.shape, dim=dim):
result, = self.node.f(input_tensor, dim)
# Ensure return value is a tensor
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
# Ensure result has correct shape
self.assertEqual(result.shape, torch.Size(expected_shape), f"Expected shape {expected_shape}, got {result.shape}")
# Ensure result has the same dtype
self.assertEqual(result.dtype, input_tensor.dtype, "Dtype mismatch after unsqueeze operation")
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/pt_view_test.py
+61
View File
@@ -0,0 +1,61 @@
import os
import sys
import unittest
import torch
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
sys.path.append(MODULE_ROOT)
from pytorch_wrapper.pt_view import PtView # Importing the PtView class
class TestPtView(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.node = PtView()
# Define test cases: (input_tensor, new_shape, expected_shape)
self.test_cases = [
# Case 1: Reshape a 2D tensor
(torch.ones(2, 3), "[3, 2]", (3, 2)), # (2,3) -> (3,2)
# Case 2: Reshape a 3D tensor
(torch.ones(2, 3, 4), "[6, 4]", (6, 4)), # (2,3,4) -> (6,4)
# Case 3: Using -1 for automatic inference
(torch.ones(2, 3, 4), "[2, -1]", (2, 12)), # (2,3,4) -> (2,12)
(torch.ones(6, 4), "[-1, 4]", (6, 4)), # (6,4) -> (6,4)
# Case 4: Flattening a multi-dimensional tensor
(torch.ones(2, 3, 4), "[-1]", (24,)), # (2,3,4) -> (24,)
# Case 5: Expanding dimensions
(torch.ones(6,), "[2, 3]", (2, 3)), # (6,) -> (2,3)
# Case 6: Reshaping a 4D tensor
(torch.ones(2, 3, 4, 5), "[6, -1, 5]", (6, 4, 5)), # (2,3,4,5) -> (6,4,5)
# Case 7: Edge case with 1D tensor
(torch.ones(10,), "[10, 1]", (10, 1)), # (10,) -> (10,1)
]
def test_view(self):
"""Test tensor reshape operation with various shape orders."""
for input_tensor, new_shape, expected_shape in self.test_cases:
with self.subTest(input_shape=input_tensor.shape, new_shape=new_shape):
result, = self.node.f(input_tensor, new_shape)
# Ensure return value is a tensor
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
# Ensure result has correct shape
self.assertEqual(result.shape, torch.Size(expected_shape), f"Expected shape {expected_shape}, got {result.shape}")
# Ensure result has the same dtype
self.assertEqual(result.dtype, input_tensor.dtype, "Dtype mismatch after view operation")
if __name__ == "__main__":
unittest.main()
tests/pytorch_wrapper/pt_zeros_test.py
+44
View File
@@ -0,0 +1,44 @@
import os
import sys
import unittest
import torch
import ast
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
sys.path.append(MODULE_ROOT)
from pytorch_wrapper.pt_zeros import PtZeros # Importing the PtZeros class
from pytorch_wrapper.utils import DTYPE_MAPPING # Importing the dtype mapping
class TestPtZeros(unittest.TestCase):
def setUp(self):
"""Set up test instance."""
self.node = PtZeros()
# Tensor size to test
self.size_str = "[2, 3]" # Using a small 2D tensor size for all tests
self.expected_shape = (2, 3)
def test_zeros_tensor_all_dtypes(self):
"""Test tensor creation with zeros for all supported data types."""
for data_type, torch_dtype in DTYPE_MAPPING.items():
with self.subTest(data_type=data_type):
result, = self.node.f(self.size_str, data_type)
# Ensure return value is a tensor
self.assertTrue(isinstance(result, torch.Tensor), "Output is not a tensor")
# Ensure result has correct shape
self.assertEqual(result.shape, torch.Size(self.expected_shape), f"Expected shape {self.expected_shape}, got {result.shape}")
# Ensure result has correct dtype
self.assertEqual(result.dtype, torch_dtype, f"Expected dtype {torch_dtype}, got {result.dtype}")
# Ensure all elements are zeros
expected = torch.zeros(self.expected_shape, dtype=torch_dtype)
self.assertTrue(torch.equal(result, expected), f"Tensor values do not match expected zeros tensor for dtype {data_type}")
if __name__ == "__main__":
unittest.main()