22 lines
630 B
Python
22 lines
630 B
Python
# Copyright (c) ONNX Project Contributors
|
|
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
|
|
from onnx.reference.op_run import OpRun
|
|
|
|
|
|
class Einsum(OpRun):
|
|
def _run(self, *args, equation=None):
|
|
if not isinstance(equation, str):
|
|
raise TypeError(f"equation must be string but is {type(equation)!r}.")
|
|
equation = equation.strip()
|
|
if not equation:
|
|
raise TypeError("equation is empty.")
|
|
try:
|
|
return (np.einsum(equation, *args, optimize=True),)
|
|
except TypeError:
|
|
return (np.einsum(equation, *args),)
|