da4ml 0.1.2__py3-none-any.whl → 0.2.0__py3-none-any.whl

da4ml/trace/fixed_variable_array.py

@@ -0,0 +1,177 @@
+from typing import Any
+
+import numpy as np
+from numpy.typing import NDArray
+
+from ..cmvm import solve
+from .fixed_variable import FixedVariable, HWConfig, QInterval
+
+
+class FixedVariableArray:
+    def __init__(
+        self,
+        vars: NDArray,
+        solver_options: dict[str, Any] | None = None,
+    ):
+        self._vars = np.array(vars)
+        self.solver_options = solver_options
+
+    @classmethod
+    def from_lhs(
+        cls,
+        low: NDArray[np.floating],
+        high: NDArray[np.floating],
+        step: NDArray[np.floating],
+        hwconf: HWConfig,
+        latency: np.ndarray | float = 0.0,
+        solver_options: dict[str, Any] | None = None,
+    ):
+        shape = low.shape
+        assert shape == high.shape == step.shape
+
+        low, high, step = low.ravel(), high.ravel(), step.ravel()
+        latency = np.full_like(low, latency) if isinstance(latency, (int, float)) else latency.ravel()
+
+        vars = []
+        for i, (l, h, s, lat) in enumerate(zip(low, high, step, latency)):
+            var = FixedVariable(
+                low=float(l),
+                high=float(h),
+                step=float(s),
+                hwconf=hwconf,
+                latency=float(
+                    lat,
+                ),
+            )
+            vars.append(var)
+        vars = np.array(vars).reshape(shape)
+        return cls(vars, solver_options)
+
+    __array_priority__ = 100
+
+    @classmethod
+    def from_kif(
+        cls,
+        k: NDArray[np.bool_ | np.integer],
+        i: NDArray[np.integer],
+        f: NDArray[np.integer],
+        hwconf: HWConfig,
+        latency: NDArray[np.floating] | float = 0.0,
+        solver_options: dict[str, Any] | None = None,
+    ):
+        step = 2.0**-f
+        _high = 2.0**i
+        high, low = _high - step, -_high * k
+        return cls.from_lhs(low, high, step, hwconf, latency, solver_options)
+
+    def __matmul__(self, other):
+        assert isinstance(other, np.ndarray)
+        kwargs = (self.solver_options or {}).copy()
+        shape0, shape1 = self.shape, other.shape
+        assert shape0[-1] == shape1[0], f'Matrix shapes do not match: {shape0} @ {shape1}'
+        c = shape1[0]
+        out_shape = shape0[:-1] + shape1[1:]
+        mat0, mat1 = self.reshape((-1, c)), other.reshape((c, -1))
+        r = []
+        for i in range(mat0.shape[0]):
+            vec = mat0[i]
+            qintervals = tuple([QInterval(float(v.low), float(v.high), float(v.step)) for v in vec._vars])
+            latencies = tuple([float(v.latency) for v in vec._vars])
+            hwconf = self._vars.ravel()[0].hwconf
+            kwargs.update(adder_size=hwconf.adder_size, carry_size=hwconf.carry_size)
+            _mat = np.ascontiguousarray(mat1.astype(np.float32))
+            sol = solve(_mat, qintervals=qintervals, latencies=latencies, **kwargs)
+            _r = sol(vec._vars)
+            r.append(_r)
+        r = np.array(r).reshape(out_shape)
+        return FixedVariableArray(r, self.solver_options)
+
+    def __rmatmul__(self, other):
+        mat1 = np.moveaxis(other, -1, 0)
+        mat0 = np.moveaxis(self._vars, 0, -1)
+        ndim0, ndim1 = mat0.ndim, mat1.ndim
+        r = FixedVariableArray(mat0, self.solver_options) @ mat1
+
+        _axes = tuple(range(0, ndim0 + ndim1 - 2))
+        axes = _axes[ndim0 - 1 :] + _axes[: ndim0 - 1]
+        return r.transpose(axes)
+
+    def __getitem__(self, *item):
+        vars = self._vars[*item]
+        if isinstance(vars, np.ndarray):
+            return FixedVariableArray(vars, self.solver_options)
+        else:
+            return vars
+
+    def __len__(self):
+        return len(self._vars)
+
+    @property
+    def shape(self):
+        return self._vars.shape
+
+    def __add__(self, other):
+        return FixedVariableArray(self._vars + other, self.solver_options)
+
+    def __sub__(self, other):
+        return FixedVariableArray(self._vars - other, self.solver_options)
+
+    def __mul__(self, other):
+        return FixedVariableArray(self._vars * other, self.solver_options)
+
+    def __truediv__(self, other):
+        return FixedVariableArray(self._vars * (1 / other), self.solver_options)
+
+    def __radd__(self, other):
+        return self + other
+
+    def __neg__(self):
+        return FixedVariableArray(-self._vars, self.solver_options)
+
+    def __repr__(self):
+        shape = self._vars.shape
+        hwconf_str = str(self._vars.ravel()[0].hwconf)[8:]
+        max_lat = max(v.latency for v in self._vars.ravel())
+        return f'FixedVariableArray(shape={shape}, hwconf={hwconf_str}, latency={max_lat})'
+
+    def relu(self, i: NDArray[np.integer] | None = None, f: NDArray[np.integer] | None = None, round_mode: str = 'TRN'):
+        shape = self._vars.shape
+        i = np.broadcast_to(i, shape) if i is not None else np.full(shape, None)
+        f = np.broadcast_to(f, shape) if f is not None else np.full(shape, None)
+        ret = []
+        for v, i, f in zip(self._vars.ravel(), i.ravel(), f.ravel()):
+            ret.append(v.relu(i=i, f=f, round_mode=round_mode))
+        return FixedVariableArray(np.array(ret).reshape(shape), self.solver_options)
+
+    def quantize(
+        self,
+        k: NDArray[np.integer] | None = None,
+        i: NDArray[np.integer] | None = None,
+        f: NDArray[np.integer] | None = None,
+        overflow_mode: str = 'WRAP',
+        round_mode: str = 'TRN',
+    ):
+        shape = self._vars.shape
+        k = np.broadcast_to(k, shape) if k is not None else np.full(shape, None)
+        i = np.broadcast_to(i, shape) if i is not None else np.full(shape, None)
+        f = np.broadcast_to(f, shape) if f is not None else np.full(shape, None)
+        ret = []
+        for v, k, i, f in zip(self._vars.ravel(), k.ravel(), i.ravel(), f.ravel()):
+            ret.append(v.quantize(k=k, i=i, f=f, overflow_mode=overflow_mode, round_mode=round_mode))
+        return FixedVariableArray(np.array(ret).reshape(shape), self.solver_options)
+
+    def flatten(self):
+        return FixedVariableArray(self._vars.flatten(), self.solver_options)
+
+    def reshape(self, shape):
+        return FixedVariableArray(self._vars.reshape(shape), self.solver_options)
+
+    def transpose(self, axes=None):
+        return FixedVariableArray(self._vars.transpose(axes), self.solver_options)
+
+    def ravel(self):
+        return FixedVariableArray(self._vars.ravel(), self.solver_options)
+
+    @property
+    def dtype(self):
+        return self._vars.dtype

da4ml/trace/ops/__init__.py

@@ -0,0 +1,55 @@
+from typing import TypeVar
+
+import numpy as np
+from numpy.typing import NDArray
+
+from ..fixed_variable_array import FixedVariable, FixedVariableArray
+from .conv_utils import conv
+from .einsum_utils import einsum
+
+T = TypeVar('T', FixedVariableArray, NDArray[np.floating], list[FixedVariable])
+
+
+def relu(x: T, i: NDArray[np.integer] | None = None, f: NDArray[np.integer] | None = None, round_mode: str = 'TRN') -> T:
+    if isinstance(x, FixedVariableArray):
+        return x.relu(i=i, f=f, round_mode=round_mode)
+    elif isinstance(x, list):
+        return [xx.relu(i=ii, f=ff, round_mode=round_mode) for xx, ii, ff in zip(x, i, f)]  # type: ignore
+    else:
+        x = np.maximum(x, 0)
+        if f is not None:
+            if round_mode.upper() == 'RND':
+                x += 2.0 ** (-f - 1)
+            sf = 2.0**f
+            x = np.floor(x * sf) / sf
+        if i is not None:
+            x = x % 2.0**i
+        return x
+
+
+def quantize(
+    x: T,
+    k: NDArray[np.integer],
+    i: NDArray[np.integer],
+    f: NDArray[np.integer],
+    overflow_mode: str = 'WRAP',
+    round_mode: str = 'TRN',
+) -> T:
+    assert overflow_mode.upper() == 'WRAP', 'Only WRAP overflow mode is supported'
+    if isinstance(x, FixedVariableArray):
+        return x.quantize(k=k, i=i, f=f, overflow_mode=overflow_mode, round_mode=round_mode)
+    else:
+        if round_mode.upper() == 'RND':
+            x += 2.0 ** (-f - 1)
+        b = k + i + f
+        bias = 2.0 ** (b - 1) * k
+        eps = 2.0**-f
+        return eps * ((np.floor(x / eps) + bias) % 2.0**b - bias)  # type: ignore
+
+
+__all__ = [
+    'conv',
+    'einsum',
+    'relu',
+    'quantize',
+]

da4ml/trace/ops/conv_utils.py

@@ -0,0 +1,104 @@
+from collections.abc import Sequence
+from typing import TypeVar
+
+import numpy as np
+from numpy.typing import NDArray
+
+from ..fixed_variable_array import FixedVariableArray
+
+
+def r_im2col(kernel_size: Sequence[int], arr: np.ndarray, buffer: np.ndarray, axis: int):
+    w = kernel_size[0]
+    if len(kernel_size) == 3:  # 1D
+        for i in range(arr.shape[axis] - w + 1):
+            patch = np.take(arr, range(i, i + w), axis=axis)
+            buffer[i] = patch.flatten()
+    else:  # 2D+
+        for i in range(arr.shape[axis] - w + 1):
+            patch = arr[i : i + w]
+            r_im2col(kernel_size[1:], patch, buffer[i], axis + 1)
+
+
+def _im2col(kernel_size: Sequence[int], arr: np.ndarray):
+    if len(kernel_size) < 3:
+        return arr
+    shape = [inp_d - ker_d + 1 for inp_d, ker_d in zip(arr.shape, kernel_size[:-2])]
+    shape.append(np.prod(kernel_size[:-1]))  # type: ignore
+    buf = np.empty(shape, dtype=arr.dtype)
+    r_im2col(kernel_size, arr, buf, 0)
+    return buf
+
+
+def stride_arr(stride: int | tuple[int, ...], arr: np.ndarray):
+    ndim = arr.ndim
+    if isinstance(stride, int):
+        stride = (stride,) * (ndim - 1)
+    assert len(stride) == ndim - 1, f'Invalid stride {stride} for array with {ndim} dimensions'
+
+    _idx = tuple(slice(None, None, st) for st in stride)
+    return arr[*_idx]
+
+
+T = TypeVar('T', FixedVariableArray, NDArray[np.integer | np.floating])
+
+
+def conv(
+    x: T,
+    kernel: NDArray[np.integer | np.floating],
+    bias: NDArray[np.integer | np.floating] | None = None,
+    strides: int | tuple[int, ...] = 1,
+    padding: tuple[tuple[int, int], ...] | str = 'VALID',
+    format: str = 'channels_last',
+):
+    if isinstance(x, FixedVariableArray):
+        solver_options = x.solver_options
+        data = x._vars
+        is_symbolic = True
+    else:
+        solver_options = None
+        data = x
+        is_symbolic = False
+
+    ndim = data.ndim
+    ch_in, ch_out = kernel.shape[-2:]
+    _ch_in = data.shape[-1]
+    assert ch_in == _ch_in, f'Invalid input shape {data.shape} for kernel {kernel.shape}'
+    assert kernel.ndim == ndim + 1
+
+    assert format in ('channels_last', 'channels_first'), f'Invalid format {format}'
+
+    if isinstance(strides, int):
+        strides = (strides,) * (ndim - 1)
+    assert len(strides) == ndim - 1, f'Invalid stride {strides} for array with {ndim} dimensions'
+
+    if isinstance(padding, str):
+        padding = padding.upper()
+        if padding == 'VALID':
+            padding = ((0, 0),) * (ndim - 1)
+        elif padding == 'SAME':
+            _padding = []
+            for i in range(ndim - 1):
+                pad0 = kernel.shape[i] // 2
+                pad1 = kernel.shape[i] - pad0 - 1
+                _padding.append((pad1, pad0))
+            padding = tuple(_padding)
+        else:
+            raise ValueError(f'Invalid padding {padding}')
+    assert len(padding) == ndim - 1, f'Invalid padding {padding} for array with {ndim} dimensions'
+    assert all(len(p) == 2 for p in padding), f'Invalid padding {padding} for array with {ndim} dimensions'
+
+    data = np.pad(data, padding + ((0, 0),), mode='constant', constant_values=0.0)
+    data = _im2col(kernel.shape, data)
+    if is_symbolic:
+        _data = FixedVariableArray(data, solver_options) @ kernel.reshape(-1, ch_out)
+        data = _data._vars
+    else:
+        data = data @ kernel.reshape(-1, ch_out)
+    data = stride_arr(strides, data)
+    if bias is not None:
+        data = data + bias
+    if format == 'channels_first':
+        data = np.moveaxis(data, -1, 1)
+    if solver_options is not None:
+        return FixedVariableArray(data, solver_options)
+    return data

da4ml/trace/ops/einsum_utils.py

@@ -0,0 +1,299 @@
+from math import prod
+from typing import TypedDict, overload
+
+import numpy as np
+from numpy.typing import NDArray
+
+from ..fixed_variable_array import FixedVariableArray
+
+
+class EinsumRecipe(TypedDict):
+    direct_sum_axis: tuple[tuple[int, ...], tuple[int, ...]]
+    in_transpose_idxs: tuple[tuple[int, ...], tuple[int, ...]]
+    L0: int
+    L1: int
+    I: int
+    C: int
+    out_interpert_shape: tuple[int, ...]
+    out_transpose_idxs: tuple[int, ...]
+
+
+def _validate_einsum_expr(fn: str, shape0: tuple[int, ...], shape1: tuple[int, ...]):
+    """Validate, resolve broadcasting, and compute output shape for einsum string
+
+    Parameters
+    ----------
+    fn : str
+        einsum string, e.g. 'ij,jk->ik'
+    shape0 : tuple[int,...]
+        shape of input0
+    shape1 : tuple[int,...]
+        shape of input1
+
+    Returns
+    -------
+    tuple[str, tuple[int,...]]
+        einsum string w/o broadcasting, and output shape
+
+    Raises
+    ------
+    ValueError
+        If the einsum string is invalid, or if it is incompatible with the input shapes
+    """
+    inp, out = map(str.strip, fn.split('->'))
+    in0, in1 = map(str.strip, inp.split(','))
+    alphabets = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'
+    s_alphabets = set(alphabets)
+
+    # Invalid characters
+    if not (s_alphabets >= set(in0.replace('...', '') + in1.replace('...', '') + out.replace('...', ''))):
+        raise ValueError(f"einsum string {fn} is invalid: subscripts should be in [a-zA-Z] and '...' only")
+
+    in0 = in0.replace('...', '0')
+    in1 = in1.replace('...', '0')
+    out = out.replace('...', '0')
+    ax_in0, ax_in1, ax_out = list(in0), list(in1), list(out)
+    sax_in0, sax_in1, sax_out = set(ax_in0), set(ax_in1), set(ax_out)
+    free_indices = ''.join(sorted(s_alphabets - sax_in0 - sax_in1 - sax_out))
+
+    # Repeated indices
+    if len(sax_in0) != len(ax_in0):
+        for a in in0:
+            if in0.count(a) == 1:
+                continue
+            a = a if a != '0' else '...'
+            raise ValueError(f"einsum string {fn} is invalid: input0 subscripts includes '{a}' multiple times")
+    if len(sax_in1) != len(ax_in1):
+        for a in in1:
+            if in1.count(a) == 1:
+                continue
+            a = a if a != '0' else '...'
+            raise ValueError(f"einsum string {fn} is invalid: input1 subscripts includes '{a}' multiple times")
+    if len(sax_out) != len(ax_out):
+        for a in out:
+            if out.count(a) == 1:
+                continue
+            a = a if a != '0' else '...'
+            raise ValueError(f"einsum string {fn} is invalid: output subscripts includes '{a}' multiple times")
+
+    # Invalid broadcasting
+    if '0' in sax_in0 or '0' in sax_in1 or '0' in sax_out:
+        if '0' not in sax_out:
+            raise ValueError(f'einsum string {fn} is invalid: output does not allow broadcasting, but inputs do')
+        if '0' not in sax_in0 and '0' not in sax_in1:
+            raise ValueError(f'einsum string {fn} is invalid: output allows broadcasting, but inputs do not')
+
+    # Output index out of nowhere
+    if remaining := sax_out - sax_in0 - sax_in1:
+        raise ValueError(f'einsum string {fn} is invalid: output subscripts {remaining} not found in inputs')
+
+    _common_in = sax_in0 & sax_in1
+
+    if '0' in sax_in0 and '0' in sax_in1:
+        # Simultaneous axes expansion in both inputs
+        n_boardcast0 = len(shape0) - len(sax_in0) + 1
+        n_boardcast1 = len(shape1) - len(sax_in1) + 1
+        assert n_boardcast0 == n_boardcast1, f'... expands to {n_boardcast0} and {n_boardcast1}-axis in input0 and input1.'
+        # Replace expansion indices with free indices
+        in0 = in0.replace('0', free_indices[:n_boardcast0])
+        in1 = in1.replace('0', free_indices[:n_boardcast1])
+        out = out.replace('0', free_indices[:n_boardcast0])
+        ax_in0, ax_in1, ax_out = list(in0), list(in1), list(out)
+        _common_in = set(ax_in0) & set(ax_in1)
+
+    else:
+        # Axes expansion in input0 or input1 only
+        if '0' in sax_in0:
+            if len(sax_in0) - 1 > len(shape0):
+                raise ValueError(f'Input0 requires at least {len(sax_in0)-1} dimensions, but only {len(shape0)} given')
+            # Replace auto expansion indices with free indices
+            n_broadcast = len(shape0) - len(sax_in0) + 1
+            in0 = in0.replace('0', free_indices[:n_broadcast])
+            out = out.replace('0', free_indices[:n_broadcast])
+            ax_in0 = list(in0)
+            ax_out = list(out)
+        else:
+            if len(sax_in0) != len(shape0):
+                raise ValueError(f'Input0 requires {len(sax_in0)} dimensions, but {len(shape0)} is given')
+
+        if '0' in sax_in1:
+            if len(sax_in1) - 1 > len(shape1):
+                raise ValueError(f'Input1 requires at least {len(sax_in1)-1} dimensions, but only {len(shape1)} given')
+            # Replace expansion indices with free indices
+            n_broadcast = len(shape1) - len(sax_in1) + 1
+            in1 = in1.replace('0', free_indices[:n_broadcast])
+            out = out.replace('0', free_indices[:n_broadcast])
+            ax_in1 = list(in1)
+            ax_out = list(out)
+        else:
+            if len(sax_in1) != len(shape1):
+                raise ValueError(f'Input1 requires {len(sax_in1)} dimensions, but {len(shape1)} is given')
+
+    # Input dimension mismatch
+    for a in _common_in:
+        ax_0 = ax_in0.index(a)
+        ax_1 = ax_in1.index(a)
+        if shape0[ax_0] != shape1[ax_1]:
+            raise ValueError(f"Input dimension size mismatches for common subscript '{a}': {shape0[ax_0]} and {shape1[ax_1]}")
+
+    out_shape = tuple(shape0[ax_in0.index(a)] if a in ax_in0 else shape1[ax_in1.index(a)] for a in ax_out)
+    return f'{in0},{in1}->{out}', out_shape
+
+
+def parse_einsum(fn: str, input_shape0: tuple[int, ...], input_shape1: tuple[int, ...]) -> EinsumRecipe:
+    """Parse einsum operation on two input arrays, return a recipe for execution
+
+    Parameters
+    ----------
+    fn : str
+        einsum string, e.g. 'ij,jk->ik'
+    input : np.ndarray
+        input0, the first input array
+    input1 : np.ndarray
+        input1, the second input array
+
+    Returns
+    -------
+    EinsumRecipe
+        einsum recipe; executed by _exec_einsum
+    """
+
+    fn, _ = _validate_einsum_expr(fn, input_shape0, input_shape1)
+
+    _in, _out = fn.split('->')
+    _in0, _in1 = _in.split(',')
+
+    in0, in1, out = list(_in0), list(_in1), list(_out)
+    s_in0, s_in1, s_out = set(in0), set(in1), set(out)
+    _common = s_in0 & s_in1
+    _contract = _common - s_out
+    _inplace = _common & s_out
+    contract = sorted(_contract, key=lambda x: in1.index(x))
+    inplace = sorted(_inplace, key=lambda x: in1.index(x))
+    invariant0 = sorted((s_out - _common) & s_in0, key=lambda x: in0.index(x))
+    invariant1 = sorted((s_out - _common) & s_in1, key=lambda x: in1.index(x))
+    direct_sum0 = s_in0 - s_out - _common
+    direct_sum1 = s_in1 - s_out - _common
+    direct_sum_axis = (
+        tuple(sorted(in0.index(x) for x in direct_sum0)),
+        tuple(sorted(in1.index(x) for x in direct_sum1)),
+    )
+
+    contract_idxs = tuple(map(in0.index, contract)), tuple(map(in1.index, contract))
+    inplace_idxs = tuple(map(in0.index, inplace)), tuple(map(in1.index, inplace))
+    invariant_idxs = tuple(map(in0.index, invariant0)), tuple(map(in1.index, invariant1))
+
+    inplace_shape = tuple(input_shape0[i] for i in inplace_idxs[0])
+    inplace_size = prod(inplace_shape)
+    contract_size = prod(input_shape0[i] for i in contract_idxs[0])
+    invariant_shape0 = tuple(input_shape0[i] for i in invariant_idxs[0])
+    invariant_shape1 = tuple(input_shape1[i] for i in invariant_idxs[1])
+    invariant_size0, invariant_size1 = prod(invariant_shape0), prod(invariant_shape1)
+
+    transpose_idx0 = inplace_idxs[0] + invariant_idxs[0] + contract_idxs[0]
+    transpose_idx1 = inplace_idxs[1] + invariant_idxs[1] + contract_idxs[1]
+
+    out_shape_pretranspose = inplace_shape + invariant_shape0 + invariant_shape1
+    _out_transpose_idx = np.argsort(tuple(map(out.index, inplace + invariant0 + invariant1)))
+    out_transpose_idx = tuple(int(i) for i in _out_transpose_idx)
+
+    return EinsumRecipe(
+        direct_sum_axis=direct_sum_axis,
+        in_transpose_idxs=(transpose_idx0, transpose_idx1),
+        out_interpert_shape=out_shape_pretranspose,
+        out_transpose_idxs=out_transpose_idx,
+        L0=invariant_size0,
+        L1=invariant_size1,
+        I=inplace_size,
+        C=contract_size,
+    )
+
+
+def _exec_einsum(recipe: EinsumRecipe, input0: np.ndarray, input1: np.ndarray) -> np.ndarray:
+    """Execute einsum operation on two input arrays
+
+    Parameters
+    ----------
+    recipe : EinsumRecipe
+        einsum recipe
+    input0 : np.ndarray
+        input0, the first input array
+    input1 : np.ndarray
+        input1, the second input array
+
+    Returns
+    -------
+    np.ndarray
+        output array
+    """
+    sum_axis0, sum_axis1 = recipe['direct_sum_axis']
+    if sum_axis0:
+        input0 = np.sum(input0, axis=sum_axis0)
+    if sum_axis1:
+        input1 = np.sum(input1, axis=sum_axis1)
+    input0 = input0.transpose(recipe['in_transpose_idxs'][0]).ravel()
+    input1 = input1.transpose(recipe['in_transpose_idxs'][1]).ravel()
+    out_dtype = object if input0.dtype == object or input1.dtype == object else np.float64
+    output = np.zeros(recipe['L0'] * recipe['L1'] * recipe['I'], dtype=out_dtype)
+
+    L0, L1, I, C = recipe['L0'], recipe['L1'], recipe['I'], recipe['C']
+
+    for l0 in range(L0):
+        for i in range(I):
+            A = input1[i * L1 * C : (i + 1) * L1 * C].reshape((L1, C))
+            B = input0[(i * L0 + l0) * C : (i * L0 + l0 + 1) * C]
+            output[(i * L0 + l0) * L1 : (i * L0 + l0 + 1) * L1] = A @ B
+
+    return output.reshape(recipe['out_interpert_shape']).transpose(recipe['out_transpose_idxs'])
+
+
+def _einsum(fn: str, input0, input1) -> np.ndarray:
+    """Execute einsum operation on two input arrays.
+
+    WARNING: Order of multiplication is reversed -- watchout if you are using non-commutative operators
+
+    Parameters
+    ----------
+    fn : str
+        einsum string, e.g. 'ij,jk->ik'
+    input : np.ndarray
+        input0, the first input array
+    input1 : np.ndarray
+        input1, the second input array
+
+    Returns
+    -------
+    np.ndarray
+        output array
+    """
+    recipe = parse_einsum(fn, input0.shape, input1.shape)
+    return _exec_einsum(recipe, input0, input1)
+
+
+@overload
+def einsum(fn: str, input0: FixedVariableArray, input1: NDArray[np.integer | np.floating]) -> FixedVariableArray: ...
+
+
+@overload
+def einsum(fn: str, input0: NDArray[np.integer | np.floating], input1: FixedVariableArray) -> FixedVariableArray: ...
+
+
+@overload
+def einsum(
+    fn: str, input0: NDArray[np.integer | np.floating], input1: NDArray[np.integer | np.floating]
+) -> NDArray[np.integer | np.floating]: ...
+
+
+def einsum(fn: str, input0, input1):
+    fg0 = isinstance(input0, FixedVariableArray)
+    fg1 = isinstance(input1, FixedVariableArray)
+    if fg0 and fg1:
+        raise ValueError('Einsum does not support two FixedVariableArray inputs')
+
+    r = _einsum(fn, input0, input1)
+    if fg0:
+        return FixedVariableArray(r, input0.solver_options)
+    elif fg1:
+        return FixedVariableArray(r, input1.solver_options)
+    else:
+        return r