da4ml 0.3.1__py3-none-any.whl → 0.3.3__py3-none-any.whl

da4ml/_version.py

@@ -1,7 +1,14 @@
 # file generated by setuptools-scm
 # don't change, don't track in version control
 
-__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
 
 TYPE_CHECKING = False
 if TYPE_CHECKING:
@@ -9,13 +16,19 @@ if TYPE_CHECKING:
     from typing import Union
 
     VERSION_TUPLE = Tuple[Union[int, str], ...]
+    COMMIT_ID = Union[str, None]
 else:
     VERSION_TUPLE = object
+    COMMIT_ID = object
 
 version: str
 __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
+commit_id: COMMIT_ID
+__commit_id__: COMMIT_ID
 
-__version__ = version = '0.3.1'
-__version_tuple__ = version_tuple = (0, 3, 1)
+__version__ = version = '0.3.3'
+__version_tuple__ = version_tuple = (0, 3, 3)
+
+__commit_id__ = commit_id = None

da4ml/cmvm/types.py

@@ -321,7 +321,7 @@ class Solution(NamedTuple):
                 case 4:  # const addition
                     bias = op.data * op.qint.step
                     buf[i] = buf[op.id0] + bias
-                case 5:
+                case 5:  # const definition
                     buf[i] = op.data * op.qint.step  # const definition
                 case 6 | -6:  # MSB Mux
                     id_c = op.data & 0xFFFFFFFF
@@ -340,6 +340,9 @@ class Solution(NamedTuple):
                         else:
                             _k, _i, _f = _minimal_kif(qint_k)
                             buf[i] = v0 if k >= 2.0 ** (_i - 1) else v1 * 2.0**shift
+                case 7:
+                    v0, v1 = buf[op.id0], buf[op.id1]
+                    buf[i] = v0 * v1
                 case _:
                     raise ValueError(f'Unknown opcode {op.opcode} in {op}')
 
@@ -370,6 +373,8 @@ class Solution(NamedTuple):
                     case 6 | -6:
                         _sign = '-' if op.opcode == -6 else ''
                         op_str = f'msb(buf[{op.data}]) ? buf[{op.id0}] : {_sign}buf[{op.id1}]'
+                    case 7:
+                        op_str = f'buf[{op.id0}] * buf[{op.id1}]'
                     case _:
                         raise ValueError(f'Unknown opcode {op.opcode} in {op}')
 
@@ -436,7 +441,12 @@ class Solution(NamedTuple):
     @property
     def inp_qint(self):
         """Quantization intervals of the input elements."""
-        return [op.qint for op in self.ops if op.opcode == -1]
+        qints = [QInterval(0.0, 0.0, 1.0) for _ in range(self.shape[0])]
+        for op in self.ops:
+            if op.opcode != -1:
+                continue
+            qints[op.id0] = op.qint
+        return qints
 
     def save(self, path: str | Path):
         """Save the solution to a file."""

da4ml/codegen/cpp/cpp_codegen.py

@@ -46,7 +46,7 @@ def ssa_gen(sol: Solution, print_latency: bool, typestr_fn: Callable[[bool | int
         match op.opcode:
             case -1:
                 # Input marker
-                val = f'inp[{ops[op.id0].id0}]'
+                val = f'inp[{op.id0}]'
             case 0 | 1:
                 # Common a+/-b<<shift op
                 ref1 = f'bit_shift<{op.data}>(v{op.id1})' if op.data != 0 else f'v{op.id1}'
@@ -86,7 +86,10 @@ def ssa_gen(sol: Solution, print_latency: bool, typestr_fn: Callable[[bool | int
                 sign = '-' if op.opcode == -6 else ''
                 ref1 = f'v{op.id1}' if shift == 0 else f'bit_shift<{shift}>(v{op.id1})'
                 val = f'{ref_k} ? {_type}({ref0}) : {_type}({sign}{ref1})'
-
+            case 7:
+                # Multiplication
+                ref1 = f'v{op.id1}'
+                val = f'{ref0} * {ref1}'
             case _:
                 raise ValueError(f'Unsupported opcode: {op.opcode}')
 

da4ml/codegen/cpp/source/binder_util.hh

@@ -9,31 +9,26 @@ constexpr bool _openmp = true;
 constexpr bool _openmp = false;
 #endif
 
-template <typename CONFIG_T, typename T> void _inference(T *c_inp, T *c_out, size_t n_samples)
-{
+template <typename CONFIG_T, typename T> void _inference(T *c_inp, T *c_out, size_t n_samples) {
     typename CONFIG_T::inp_t in_fixed_buf[CONFIG_T::N_inp];
     typename CONFIG_T::out_t out_fixed_buf[CONFIG_T::N_out];
 
-    for(size_t i = 0; i < n_samples; ++i)
-        {
-            size_t offset_in = i * CONFIG_T::N_inp;
-            size_t offset_out = i * CONFIG_T::N_out;
-            for(size_t j = 0; j < CONFIG_T::N_inp; ++j)
-                {
-                    in_fixed_buf[j] = c_inp[offset_in + j];
-                }
+    for (size_t i = 0; i < n_samples; ++i) {
+        size_t offset_in = i * CONFIG_T::N_inp;
+        size_t offset_out = i * CONFIG_T::N_out;
+        for (size_t j = 0; j < CONFIG_T::N_inp; ++j) {
+            in_fixed_buf[j] = c_inp[offset_in + j];
+        }
 
-            CONFIG_T::f(in_fixed_buf, out_fixed_buf);
+        CONFIG_T::f(in_fixed_buf, out_fixed_buf);
 
-            for(size_t j = 0; j < CONFIG_T::N_out; ++j)
-                {
-                    c_out[offset_out + j] = out_fixed_buf[j];
-                }
+        for (size_t j = 0; j < CONFIG_T::N_out; ++j) {
+            c_out[offset_out + j] = out_fixed_buf[j];
         }
+    }
 }
 
-template <typename CONFIG_T, typename T> void batch_inference(T *c_inp, T *c_out, size_t n_samples)
-{
+template <typename CONFIG_T, typename T> void batch_inference(T *c_inp, T *c_out, size_t n_samples) {
 #ifdef _OPENMP
     size_t n_max_threads = omp_get_max_threads();
     size_t n_samples_per_thread = std::max<size_t>(n_samples / n_max_threads, 32);
@@ -41,15 +36,14 @@ template <typename CONFIG_T, typename T> void batch_inference(T *c_inp, T *c_out
     n_thread += (n_samples % n_samples_per_thread) ? 1 : 0;
 
 #pragma omp parallel for num_threads(n_thread) schedule(static)
-    for(size_t i = 0; i < n_thread; ++i)
-        {
-            size_t start = i * n_samples_per_thread;
-            size_t end = std::min<size_t>(start + n_samples_per_thread, n_samples);
-            size_t n_samples_this_thread = end - start;
-            size_t offset_in = start * CONFIG_T::N_inp;
-            size_t offset_out = start * CONFIG_T::N_out;
-            _inference<CONFIG_T, T>(&c_inp[offset_in], &c_out[offset_out], n_samples_this_thread);
-        }
+    for (size_t i = 0; i < n_thread; ++i) {
+        size_t start = i * n_samples_per_thread;
+        size_t end = std::min<size_t>(start + n_samples_per_thread, n_samples);
+        size_t n_samples_this_thread = end - start;
+        size_t offset_in = start * CONFIG_T::N_inp;
+        size_t offset_out = start * CONFIG_T::N_out;
+        _inference<CONFIG_T, T>(&c_inp[offset_in], &c_out[offset_out], n_samples_this_thread);
+    }
 #else
     _inference<CONFIG_T, T>(c_inp, c_out, n_samples);
 #endif

da4ml/codegen/cpp/source/vitis_bitshift.hh

@@ -1,14 +1,16 @@
 #pragma once
-#include "ap_types/ap_fixed.h"
+#include "ap_fixed.h"
 
-template <int s, int b, int i, ap_q_mode Q, ap_o_mode O, int N> ap_fixed<b, i + s> bit_shift(ap_fixed<b, i, Q, O, N> x) {
+template <int s, int b, int i, ap_q_mode Q, ap_o_mode O, int N>
+ap_fixed<b, i + s> bit_shift(ap_fixed<b, i, Q, O, N> x) {
 #pragma HLS INLINE
     ap_fixed<b, i + s> r;
     r.range() = x.range();
     return r;
 };
 
-template <int s, int b, int i, ap_q_mode Q, ap_o_mode O, int N> ap_ufixed<b, i + s> bit_shift(ap_ufixed<b, i, Q, O, N> x) {
+template <int s, int b, int i, ap_q_mode Q, ap_o_mode O, int N>
+ap_ufixed<b, i + s> bit_shift(ap_ufixed<b, i, Q, O, N> x) {
 #pragma HLS INLINE
     ap_ufixed<b, i + s> r;
     r.range() = x.range();

da4ml/codegen/verilog/comb.py

@@ -9,7 +9,7 @@ def ssa_gen(sol: Solution, neg_defined: set[int], print_latency: bool = False):
     ops = sol.ops
     kifs = list(map(_minimal_kif, (op.qint for op in ops)))
     widths = list(map(sum, kifs))
-    inp_kifs = [_minimal_kif(op.qint) for op in ops if op.opcode == -1]
+    inp_kifs = [_minimal_kif(qint) for qint in sol.inp_qint]
     inp_widths = list(map(sum, inp_kifs))
     _inp_widths = np.cumsum([0] + inp_widths)
     inp_idxs = np.stack([_inp_widths[1:] - 1, _inp_widths[:-1]], axis=1)
@@ -31,6 +31,17 @@ def ssa_gen(sol: Solution, neg_defined: set[int], print_latency: bool = False):
             case -1:  # Input marker
                 i0, i1 = inp_idxs[op.id0]
                 line = f'{_def} assign {v} = inp[{i0}:{i1}];'
+
+            case 0 | 1:  # Common a+/-b<<shift oprs
+                p0, p1 = kifs[op.id0], kifs[op.id1]  # precision -> keep_neg, integers (no sign), fractional
+
+                bw0, bw1 = widths[op.id0], widths[op.id1]  # width
+                s0, f0, s1, f1 = int(p0[0]), p0[2], int(p1[0]), p1[2]
+                shift = op.data + f0 - f1
+                v0, v1 = f'v{op.id0}[{bw0 - 1}:0]', f'v{op.id1}[{bw1 - 1}:0]'
+
+                line = f'{_def} shift_adder #({bw0}, {bw1}, {s0}, {s1}, {bw}, {shift}, {op.opcode}) op_{i} ({v0}, {v1}, {v});'
+
             case 2 | -2:  # ReLU
                 lsb_bias = kifs[op.id0][2] - kifs[i][2]
                 i0, i1 = bw + lsb_bias - 1, lsb_bias
@@ -93,16 +104,6 @@ def ssa_gen(sol: Solution, neg_defined: set[int], print_latency: bool = False):
                     num = 2**bw + num
                 line = f"{_def} assign {v} = '{bin(num)[1:]};"
 
-            case 0 | 1:  # Common a+/-b<<shift oprs
-                p0, p1 = kifs[op.id0], kifs[op.id1]  # precision -> keep_neg, integers (no sign), fractional
-
-                bw0, bw1 = widths[op.id0], widths[op.id1]  # width
-                s0, f0, s1, f1 = int(p0[0]), p0[2], int(p1[0]), p1[2]
-                shift = op.data + f0 - f1
-                v0, v1 = f'v{op.id0}[{bw0 - 1}:0]', f'v{op.id1}[{bw1 - 1}:0]'
-
-                line = f'{_def} shift_adder #({bw0}, {bw1}, {s0}, {s1}, {bw}, {shift}, {op.opcode}) op_{i} ({v0}, {v1}, {v});'
-
             case 6 | -6:  # MSB Muxing
                 k, a, b = op.data & 0xFFFFFFFF, op.id0, op.id1
                 p0, p1 = kifs[a], kifs[b]
@@ -115,6 +116,13 @@ def ssa_gen(sol: Solution, neg_defined: set[int], print_latency: bool = False):
                 vk, v0, v1 = f'v{k}[{bwk - 1}]', f'v{a}[{bw0 - 1}:0]', f'v{b}[{bw1 - 1}:0]'
 
                 line = f'{_def} mux #({bw0}, {bw1}, {s0}, {s1}, {bw}, {shift}, {inv}) op_{i} ({vk}, {v0}, {v1}, {v});'
+            case 7:  # Multiplication
+                bw0, bw1 = widths[op.id0], widths[op.id1]  # width
+                s0, s1 = int(kifs[op.id0][0]), int(kifs[op.id1][0])
+                v0, v1 = f'v{op.id0}[{bw0 - 1}:0]', f'v{op.id1}[{bw1 - 1}:0]'
+
+                line = f'{_def} multiplier #({bw0}, {bw1}, {s0}, {s1}, {bw}) op_{i} ({v0}, {v1}, {v});'
+
             case _:
                 raise ValueError(f'Unknown opcode {op.opcode} for operation {i} ({op})')
 

da4ml/codegen/verilog/source/binder_util.hh

@@ -10,7 +10,7 @@ constexpr bool _openmp = false;
 
 template <typename CONFIG_T>
 std::enable_if_t<CONFIG_T::II != 0> _inference(int32_t *c_inp, int32_t *c_out, size_t n_samples) {
-    typename CONFIG_T::dut_t *dut = new typename CONFIG_T::dut_t;
+    auto dut = std::make_unique<typename CONFIG_T::dut_t>();
 
     size_t clk_req = n_samples * CONFIG_T::II + CONFIG_T::latency + 1;
 
@@ -18,14 +18,18 @@ std::enable_if_t<CONFIG_T::II != 0> _inference(int32_t *c_inp, int32_t *c_out, s
         size_t t_out = t_inp - CONFIG_T::latency - 1;
 
         if (t_inp < n_samples * CONFIG_T::II && t_inp % CONFIG_T::II == 0) {
-            write_input<CONFIG_T::N_inp, CONFIG_T::max_inp_bw>(dut->inp, &c_inp[t_inp / CONFIG_T::II * CONFIG_T::N_inp]);
+            write_input<CONFIG_T::N_inp, CONFIG_T::max_inp_bw>(
+                dut->inp, &c_inp[t_inp / CONFIG_T::II * CONFIG_T::N_inp]
+            );
         }
 
         dut->clk = 0;
         dut->eval();
 
         if (t_inp > CONFIG_T::latency && t_out % CONFIG_T::II == 0) {
-            read_output<CONFIG_T::N_out, CONFIG_T::max_out_bw>(dut->out, &c_out[t_out / CONFIG_T::II * CONFIG_T::N_out]);
+            read_output<CONFIG_T::N_out, CONFIG_T::max_out_bw>(
+                dut->out, &c_out[t_out / CONFIG_T::II * CONFIG_T::N_out]
+            );
         }
 
         dut->clk = 1;
@@ -33,12 +37,11 @@ std::enable_if_t<CONFIG_T::II != 0> _inference(int32_t *c_inp, int32_t *c_out, s
     }
 
     dut->final();
-    delete dut;
 }
 
 template <typename CONFIG_T>
 std::enable_if_t<CONFIG_T::II == 0> _inference(int32_t *c_inp, int32_t *c_out, size_t n_samples) {
-    typename CONFIG_T::dut_t *dut = new typename CONFIG_T::dut_t;
+    auto dut = std::make_unique<typename CONFIG_T::dut_t>();
 
     for (size_t i = 0; i < n_samples; ++i) {
         write_input<CONFIG_T::N_inp, CONFIG_T::max_inp_bw>(dut->inp, &c_inp[i * CONFIG_T::N_inp]);
@@ -47,7 +50,6 @@ std::enable_if_t<CONFIG_T::II == 0> _inference(int32_t *c_inp, int32_t *c_out, s
     }
 
     dut->final();
-    delete dut;
 }
 
 template <typename CONFIG_T> void batch_inference(int32_t *c_inp, int32_t *c_out, size_t n_samples) {

da4ml/codegen/verilog/source/ioutil.hh

@@ -68,7 +68,8 @@ template <size_t bw, size_t N_out> std::vector<int32_t> bitunpack(const std::vec
 }
 
 template <size_t bits_in, typename inp_buf_t>
-std::enable_if_t<std::is_integral_v<inp_buf_t>, void> _write_input(inp_buf_t &inp_buf, const std::vector<int32_t> &input) {
+std::enable_if_t<std::is_integral_v<inp_buf_t>, void>
+_write_input(inp_buf_t &inp_buf, const std::vector<int32_t> &input) {
     assert(input.size() == (bits_in + 31) / 32);
     inp_buf = input[0] & 0xFFFFFFFF;
     if (bits_in > 32) {

da4ml/codegen/verilog/source/multiplier.v

@@ -0,0 +1,37 @@
+`timescale 1ns / 1ps
+
+
+module multiplier #(
+    parameter BW_INPUT0 = 32,
+    parameter BW_INPUT1 = 32,
+    parameter SIGNED0 = 0,
+    parameter SIGNED1 = 0,
+    parameter BW_OUT = 32
+) (
+    input [BW_INPUT0-1:0] in0,
+    input [BW_INPUT1-1:0] in1,
+    output [BW_OUT-1:0] out
+);
+
+  localparam BW_BUF = BW_INPUT0 + BW_INPUT1;
+
+  // verilator lint_off UNUSEDSIGNAL
+  wire [BW_BUF - 1:0] buffer;
+  // verilator lint_on UNUSEDSIGNAL
+
+  generate
+    if (SIGNED0 == 1 && SIGNED1 == 1) begin : signed_signed
+      assign buffer[BW_BUF-1:0] = $signed(in0) * $signed(in1);
+    end else if (SIGNED0 == 1 && SIGNED1 == 0) begin : signed_unsigned
+      assign buffer[BW_BUF-1:0] = $signed(in0) * $signed({{1'b0,in1}});
+      // assign buffer[BW_BUF-1] = in0[BW_INPUT0-1];
+    end else if (SIGNED0 == 0 && SIGNED1 == 1) begin : unsigned_signed
+      assign buffer[BW_BUF-1:0] = $signed({{1'b0,in0}}) * $signed(in1);
+      // assign buffer[BW_BUF-1] = in1[BW_INPUT1-1];
+    end else begin : unsigned_unsigned
+      assign buffer[BW_BUF-1:0] = in0 * in1;
+    end
+  endgenerate
+
+  assign out[BW_OUT-1:0] = buffer[BW_OUT-1:0];
+endmodule

da4ml/codegen/verilog/verilog_model.py

@@ -114,9 +114,8 @@ class VerilogModel:
             f.write(binder)
 
         # Common resource copy
-        shutil.copy(self.__src_root / 'verilog/source/shift_adder.v', self._path)
-        shutil.copy(self.__src_root / 'verilog/source/mux.v', self._path)
-        shutil.copy(self.__src_root / 'verilog/source/negative.v', self._path)
+        for fname in self.__src_root.glob('verilog/source/*.v'):
+            shutil.copy(fname, self._path)
         shutil.copy(self.__src_root / 'verilog/source/build_binder.mk', self._path)
         shutil.copy(self.__src_root / 'verilog/source/ioutil.hh', self._path)
         shutil.copy(self.__src_root / 'verilog/source/binder_util.hh', self._path)

da4ml/converter/__init__.py

@@ -0,0 +1,3 @@
+from .hgq2 import trace_model
+
+__all__ = ['trace_model']

da4ml/converter/hgq2/__init__.py

@@ -0,0 +1,3 @@
+from .parser import trace_model
+
+__all__ = ['trace_model']

da4ml/converter/hgq2/parser.py

@@ -1,12 +1,13 @@
 from collections.abc import Sequence
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, Literal, overload
 
 import keras
+import numpy as np
 from keras import KerasTensor, Operation
 
-from ...trace import FixedVariableArray, HWConfig
-from ...trace.fixed_variable_array import FixedVariableArrayInput
+from ...trace import FixedVariableArray, FixedVariableArrayInput, HWConfig, comb_trace
+from ...trace.fixed_variable import FixedVariable
 from .replica import _registry
 
 
@@ -20,6 +21,8 @@ class OpObj:
 
 
 def parse_model(model: keras.Model):
+    if isinstance(model, keras.Sequential):
+        model = model._functional
     operators: dict[int, list[OpObj]] = {}
     for depth, nodes in model._nodes_by_depth.items():
         _oprs = []
@@ -49,9 +52,24 @@ def replace_tensors(tensor_map: dict[KerasTensor, FixedVariableArray], obj: Any)
     return obj
 
 
+def _flatten_arr(args: Any) -> FixedVariableArray:
+    if isinstance(args, FixedVariableArray):
+        return np.ravel(args)  # type: ignore
+    if isinstance(args, FixedVariable):
+        return FixedVariableArray(np.array([args]))
+    if not isinstance(args, Sequence):
+        return None  # type: ignore
+    args = [_flatten_arr(a) for a in args]
+    args = [a for a in args if a is not None]
+    return np.concatenate(args)  # type: ignore
+
+
 def _apply_nn(
-    model: keras.Model, inputs: FixedVariableArray | Sequence[FixedVariableArray], verbose: bool = False
-) -> tuple[FixedVariableArray, ...]:
+    model: keras.Model,
+    inputs: FixedVariableArray | Sequence[FixedVariableArray],
+    verbose: bool = False,
+    dump: bool = False,
+) -> tuple[FixedVariableArray, ...] | dict[str, FixedVariableArray]:
     """
     Apply a keras model to a fixed variable array or a sequence of fixed variable arrays.
 
@@ -73,6 +91,8 @@ def _apply_nn(
     assert len(model.inputs) == len(inputs), f'Model has {len(model.inputs)} inputs, got {len(inputs)}'
     tensor_map = {keras_tensor: da_tensor for keras_tensor, da_tensor in zip(model.inputs, inputs)}
 
+    _inputs = _flatten_arr(inputs)
+
     for ops in parse_model(model):
         for op in ops:
             assert all(t in tensor_map for t in op.requires)
@@ -82,24 +102,56 @@ def _apply_nn(
                 continue
             mirror_op = _registry[op.operation.__class__](op.operation)
             if verbose:
-                print(f'Processing operation {op.operation.name} ({op.operation.__class__.__name__})')
+                print(f'Processing operation {op.operation.name} ({op.operation.__class__.__name__})', end='')
             outputs = mirror_op(*args, **kwargs)
             for keras_tensor, da_tensor in zip(op.produces, outputs):
                 tensor_map[keras_tensor] = da_tensor
+            if verbose:
+                cost = comb_trace(_inputs, _flatten_arr(outputs)).cost
+                print(f' cumcost: {cost}')
+
+    if not dump:
+        return tuple(tensor_map[keras_tensor] for keras_tensor in model.outputs)
+    else:
+        return {k.name: v for k, v in tensor_map.items()}
+
+
+@overload
+def trace_model(  # type: ignore
+    model: keras.Model,
+    hwconf: HWConfig = HWConfig(1, -1, -1),
+    solver_options: dict[str, Any] | None = None,
+    verbose: bool = False,
+    inputs: tuple[FixedVariableArray, ...] | FixedVariableArray | None = None,
+    dump: Literal[False] = False,
+) -> tuple[FixedVariableArray, FixedVariableArray]: ...
+
 
-    return tuple(tensor_map[keras_tensor] for keras_tensor in model.outputs)
+@overload
+def trace_model(  # type: ignore
+    model: keras.Model,
+    hwconf: HWConfig = HWConfig(1, -1, -1),
+    solver_options: dict[str, Any] | None = None,
+    verbose: bool = False,
+    inputs: tuple[FixedVariableArray, ...] | FixedVariableArray | None = None,
+    dump: Literal[True] = False,  # type: ignore
+) -> dict[str, FixedVariableArray]: ...
 
 
-def trace_model(
+def trace_model(  # type: ignore
     model: keras.Model,
     hwconf: HWConfig = HWConfig(1, -1, -1),
     solver_options: dict[str, Any] | None = None,
     verbose: bool = False,
     inputs: tuple[FixedVariableArray, ...] | None = None,
-) -> tuple[tuple[FixedVariableArray, ...], tuple[FixedVariableArray, ...]]:
+    dump=False,
+):
     if inputs is None:
         inputs = tuple(
             FixedVariableArrayInput(inp.shape[1:], hwconf=hwconf, solver_options=solver_options) for inp in model.inputs
         )
-    outputs = _apply_nn(model, inputs, verbose=verbose)
-    return inputs, outputs
+    outputs = _apply_nn(model, inputs, verbose=verbose, dump=dump)
+    if not dump:
+        return _flatten_arr(inputs), _flatten_arr(outputs)
+    else:
+        return {k: _flatten_arr(v) for k, v in outputs.items()}  # type: ignore