diff --git a/__init__.py b/__init__.py
index d7f7ca3..959a888 100644
--- a/__init__.py
+++ b/__init__.py
@@ -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",
diff --git a/docs/reference/node_reference.md b/docs/reference/node_reference.md
index 3286616..c127412 100644
--- a/docs/reference/node_reference.md
+++ b/docs/reference/node_reference.md
@@ -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 |
| --- | --- |
diff --git a/docs/reference/pt_full.md b/docs/reference/pt_full.md
new file mode 100644
index 0000000..df0a41b
--- /dev/null
+++ b/docs/reference/pt_full.md
@@ -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 |
+
+
+Category: PyTorch wrapper - Tensor creation
+
+ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
diff --git a/docs/reference/pt_ones.md b/docs/reference/pt_ones.md
new file mode 100644
index 0000000..5f42dc7
--- /dev/null
+++ b/docs/reference/pt_ones.md
@@ -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 |
+
+
+Category: PyTorch wrapper - Tensor creation
+
+ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
diff --git a/docs/reference/pt_permute.md b/docs/reference/pt_permute.md
new file mode 100644
index 0000000..8fa187f
--- /dev/null
+++ b/docs/reference/pt_permute.md
@@ -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 |
+
+
+Category: PyTorch wrapper - Transform
+
+ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
diff --git a/docs/reference/pt_rand.md b/docs/reference/pt_rand.md
new file mode 100644
index 0000000..34ccd58
--- /dev/null
+++ b/docs/reference/pt_rand.md
@@ -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 |
+
+
+Category: PyTorch wrapper - Tensor creation
+
+ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
diff --git a/docs/reference/pt_rand_int.md b/docs/reference/pt_rand_int.md
new file mode 100644
index 0000000..16552cd
--- /dev/null
+++ b/docs/reference/pt_rand_int.md
@@ -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 |
+
+
+Category: PyTorch wrapper - Tensor creation
+
+ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
diff --git a/docs/reference/pt_randn.md b/docs/reference/pt_randn.md
new file mode 100644
index 0000000..8c2057b
--- /dev/null
+++ b/docs/reference/pt_randn.md
@@ -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 |
+
+
+Category: PyTorch wrapper - Tensor creation
+
+ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
diff --git a/docs/reference/pt_squeeze.md b/docs/reference/pt_squeeze.md
new file mode 100644
index 0000000..b11c993
--- /dev/null
+++ b/docs/reference/pt_squeeze.md
@@ -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 |
+
+
+Category: PyTorch wrapper - Transform
+
+ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
diff --git a/docs/reference/pt_unsqueeze.md b/docs/reference/pt_unsqueeze.md
new file mode 100644
index 0000000..96bc5b1
--- /dev/null
+++ b/docs/reference/pt_unsqueeze.md
@@ -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 |
+
+
+Category: PyTorch wrapper - Transform
+
+ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
diff --git a/docs/reference/pt_view.md b/docs/reference/pt_view.md
new file mode 100644
index 0000000..d366f3e
--- /dev/null
+++ b/docs/reference/pt_view.md
@@ -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 |
+
+
+Category: PyTorch wrapper - Transform
+
+ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
diff --git a/docs/reference/pt_zeros.md b/docs/reference/pt_zeros.md
new file mode 100644
index 0000000..7019740
--- /dev/null
+++ b/docs/reference/pt_zeros.md
@@ -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 |
+
+
+Category: PyTorch wrapper - Tensor creation
+
+ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
diff --git a/modules/pytorch_wrapper/pt_full.py b/modules/pytorch_wrapper/pt_full.py
new file mode 100644
index 0000000..8ce80a8
--- /dev/null
+++ b/modules/pytorch_wrapper/pt_full.py
@@ -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,)
diff --git a/modules/pytorch_wrapper/pt_ones.py b/modules/pytorch_wrapper/pt_ones.py
new file mode 100644
index 0000000..7d6131e
--- /dev/null
+++ b/modules/pytorch_wrapper/pt_ones.py
@@ -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,)
diff --git a/modules/pytorch_wrapper/pt_permute.py b/modules/pytorch_wrapper/pt_permute.py
new file mode 100644
index 0000000..86ea0f8
--- /dev/null
+++ b/modules/pytorch_wrapper/pt_permute.py
@@ -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,)
diff --git a/modules/pytorch_wrapper/pt_rand.py b/modules/pytorch_wrapper/pt_rand.py
new file mode 100644
index 0000000..3d967ea
--- /dev/null
+++ b/modules/pytorch_wrapper/pt_rand.py
@@ -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,)
diff --git a/modules/pytorch_wrapper/pt_rand_int.py b/modules/pytorch_wrapper/pt_rand_int.py
new file mode 100644
index 0000000..071ea48
--- /dev/null
+++ b/modules/pytorch_wrapper/pt_rand_int.py
@@ -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,)
diff --git a/modules/pytorch_wrapper/pt_randn.py b/modules/pytorch_wrapper/pt_randn.py
new file mode 100644
index 0000000..9084769
--- /dev/null
+++ b/modules/pytorch_wrapper/pt_randn.py
@@ -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,)
diff --git a/modules/pytorch_wrapper/pt_squeeze.py b/modules/pytorch_wrapper/pt_squeeze.py
new file mode 100644
index 0000000..c86f3ce
--- /dev/null
+++ b/modules/pytorch_wrapper/pt_squeeze.py
@@ -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,)
diff --git a/modules/pytorch_wrapper/pt_unsqueeze.py b/modules/pytorch_wrapper/pt_unsqueeze.py
new file mode 100644
index 0000000..7505b95
--- /dev/null
+++ b/modules/pytorch_wrapper/pt_unsqueeze.py
@@ -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,)
diff --git a/modules/pytorch_wrapper/pt_view.py b/modules/pytorch_wrapper/pt_view.py
new file mode 100644
index 0000000..3d7de7d
--- /dev/null
+++ b/modules/pytorch_wrapper/pt_view.py
@@ -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,)
diff --git a/modules/pytorch_wrapper/pt_zeros.py b/modules/pytorch_wrapper/pt_zeros.py
new file mode 100644
index 0000000..0eeacc0
--- /dev/null
+++ b/modules/pytorch_wrapper/pt_zeros.py
@@ -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,)
diff --git a/modules/pytorch_wrapper/utils.py b/modules/pytorch_wrapper/utils.py
new file mode 100644
index 0000000..ad32823
--- /dev/null
+++ b/modules/pytorch_wrapper/utils.py
@@ -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
+}
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
index 961e54e..bd4169c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -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]
diff --git a/tests/pytorch_wrapper/pt_full_test.py b/tests/pytorch_wrapper/pt_full_test.py
new file mode 100644
index 0000000..77f35de
--- /dev/null
+++ b/tests/pytorch_wrapper/pt_full_test.py
@@ -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()
diff --git a/tests/pytorch_wrapper/pt_ones_test.py b/tests/pytorch_wrapper/pt_ones_test.py
new file mode 100644
index 0000000..fc893af
--- /dev/null
+++ b/tests/pytorch_wrapper/pt_ones_test.py
@@ -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()
diff --git a/tests/pytorch_wrapper/pt_permute_test.py b/tests/pytorch_wrapper/pt_permute_test.py
new file mode 100644
index 0000000..04a8e74
--- /dev/null
+++ b/tests/pytorch_wrapper/pt_permute_test.py
@@ -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()
diff --git a/tests/pytorch_wrapper/pt_rand_int_test.py b/tests/pytorch_wrapper/pt_rand_int_test.py
new file mode 100644
index 0000000..57c66e1
--- /dev/null
+++ b/tests/pytorch_wrapper/pt_rand_int_test.py
@@ -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()
diff --git a/tests/pytorch_wrapper/pt_rand_test.py b/tests/pytorch_wrapper/pt_rand_test.py
new file mode 100644
index 0000000..c3ef56c
--- /dev/null
+++ b/tests/pytorch_wrapper/pt_rand_test.py
@@ -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()
diff --git a/tests/pytorch_wrapper/pt_randn_test.py b/tests/pytorch_wrapper/pt_randn_test.py
new file mode 100644
index 0000000..9022c1d
--- /dev/null
+++ b/tests/pytorch_wrapper/pt_randn_test.py
@@ -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()
diff --git a/tests/pytorch_wrapper/pt_squeeze_test.py b/tests/pytorch_wrapper/pt_squeeze_test.py
new file mode 100644
index 0000000..a3d79f8
--- /dev/null
+++ b/tests/pytorch_wrapper/pt_squeeze_test.py
@@ -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()
diff --git a/tests/pytorch_wrapper/pt_unsqueeze_test.py b/tests/pytorch_wrapper/pt_unsqueeze_test.py
new file mode 100644
index 0000000..d9b97a7
--- /dev/null
+++ b/tests/pytorch_wrapper/pt_unsqueeze_test.py
@@ -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()
diff --git a/tests/pytorch_wrapper/pt_view_test.py b/tests/pytorch_wrapper/pt_view_test.py
new file mode 100644
index 0000000..a35dc48
--- /dev/null
+++ b/tests/pytorch_wrapper/pt_view_test.py
@@ -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()
diff --git a/tests/pytorch_wrapper/pt_zeros_test.py b/tests/pytorch_wrapper/pt_zeros_test.py
new file mode 100644
index 0000000..66cfb6d
--- /dev/null
+++ b/tests/pytorch_wrapper/pt_zeros_test.py
@@ -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()