40 lines
1.0 KiB
Python
40 lines
1.0 KiB
Python
from typing import Any, Dict, Callable, Tuple
|
|
import torch
|
|
import torch.nn.functional
|
|
import functools
|
|
|
|
log_softmax_last_axis = functools.partial(torch.nn.functional.log_softmax, dim=-1)
|
|
|
|
class PtfLogSoftmax:
|
|
"""
|
|
The log softmax activation function.
|
|
|
|
This class provides a callable that applies the log_softmax activation function
|
|
from `torch.nn.functional`.
|
|
|
|
category: PyTorch wrapper - Callable
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> Dict[str, Any]:
|
|
"""
|
|
Defines the required input types.
|
|
|
|
Returns:
|
|
Dict[str, Any]: An empty dictionary since this function requires no inputs.
|
|
"""
|
|
return {}
|
|
|
|
RETURN_TYPES: Tuple[str, ...] = ("PTCALLABLE",)
|
|
FUNCTION: str = "f"
|
|
CATEGORY: str = "Data Analysis"
|
|
|
|
def f(self) -> Tuple[Callable]:
|
|
"""
|
|
Returns the PyTorch log_softmax function as a callable.
|
|
|
|
Returns:
|
|
Tuple[Callable]: A tuple containing `logsoftmax_last_axis`.
|
|
"""
|
|
return (log_softmax_last_axis,)
|