68 lines
2.0 KiB
Python
68 lines
2.0 KiB
Python
from typing import Any, Dict
|
|
import torch
|
|
|
|
|
|
class PtEinsum:
|
|
"""
|
|
Performs Tensor operations specified in the Einstein summation equation.
|
|
|
|
Specify Einstein summation notation in the first socket.
|
|
Specify two tensors in tens_a, tens_b. If you specify the third tensor, specify in tens_c.
|
|
Otherwise, leave the socket unconnected.
|
|
|
|
Sublist format is not supported.
|
|
|
|
category: PyTorch wrapper - Matrix operations
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> Dict[str, Any]:
|
|
"""
|
|
Defines the required input types.
|
|
|
|
Returns:
|
|
Dict[str, Any]: A dictionary of required input types.
|
|
"""
|
|
return {
|
|
"required": {
|
|
"equation": ("STRING", {"default": ""}),
|
|
"tens_a": ("TENSOR", {})
|
|
},
|
|
"optional":{
|
|
"tens_b": ("TENSOR", {}),
|
|
"tens_c": ("TENSOR", {})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES: tuple = ("TENSOR",)
|
|
FUNCTION: str = "f"
|
|
CATEGORY: str = "Data Analysis"
|
|
|
|
def f(self,
|
|
equation: str,
|
|
tens_a: torch.Tensor,
|
|
tens_b: torch.Tensor = None,
|
|
tens_c: torch.Tensor = None) -> tuple:
|
|
"""
|
|
Performs Tensor operations specified in the Einstein summation equation.
|
|
|
|
Args:
|
|
equation (str): Einstein summation equation.
|
|
tens_a (torch.Tensor): A PyTorch tensor.
|
|
tens_b (torch.Tensor): A PyTorch tensor.
|
|
tens_c (torch.Tensor): A PyTorch tensor.
|
|
|
|
Returns:
|
|
tuple: A tuple containing the resultant tensor.
|
|
"""
|
|
if tens_a is None:
|
|
raise ValueError("tens_a is None")
|
|
|
|
if tens_b is None and tens_c is None:
|
|
tens_ret_val = torch.einsum(equation, tens_a)
|
|
elif tens_c is None:
|
|
tens_ret_val = torch.einsum(equation, tens_a, tens_b)
|
|
else:
|
|
tens_ret_val = torch.einsum(equation, tens_a, tens_b, tens_c)
|
|
return (tens_ret_val,)
|