62 lines
2.1 KiB
Python
62 lines
2.1 KiB
Python
import os
|
|
import sys
|
|
from typing import Any, Dict, Tuple
|
|
import torch
|
|
from .utils import DTYPE_MAPPING
|
|
|
|
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)
|
|
|
|
|
|
class PtLinspace:
|
|
"""
|
|
Creates a PyTorch tensor using `torch.linspace` with the specified start, end, and steps values.
|
|
|
|
The start and end values are parsed from string inputs and converted to a float.
|
|
The steps value is converted to an int.
|
|
|
|
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": {
|
|
"start": ("STRING", {"default":"", "multiline": False}),
|
|
"end": ("STRING", {"default":"", "multiline": False}),
|
|
"steps": ("STRING", {"default":"", "multiline": False}),
|
|
"data_type": (("float32", "float16", "bfloat16", "float64", "uint8", "int8", "int16", "int32", "int64"),)
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES: tuple = ("TENSOR",)
|
|
FUNCTION: str = "f"
|
|
CATEGORY: str = "Data Analysis"
|
|
|
|
|
|
def f(self, start: str, end: str, steps: str, data_type: str) -> Tuple[torch.Tensor]:
|
|
"""
|
|
Creates a PyTorch tensor using `torch.arange` with the given start, end, and step values.
|
|
|
|
Args:
|
|
start (str): The starting value of the sequence.
|
|
end (str): The end value of the sequence (inclusive).
|
|
steps (str): The number of steps
|
|
data_type (str): The desired PyTorch dtype for the resulting tensor.
|
|
|
|
Returns:
|
|
Tuple[torch.Tensor]: A tuple containing the generated PyTorch tensor.
|
|
"""
|
|
start_value = float(start)
|
|
end_value = float(end)
|
|
steps_value = int(steps)
|
|
tens = torch.linspace(start_value, end_value, steps_value, dtype=DTYPE_MAPPING[data_type])
|
|
return (tens,)
|