alkaid.trace.ops package

Submodules

alkaid.trace.ops.bit_oprs module

alkaid.trace.ops.bit_oprs.binary_bit_op(a: FVariable, b: FVariable, op: int, *args, **kwargs) → FVariable
alkaid.trace.ops.bit_oprs.binary_bit_op(a: float, b: float, op: int, qint0: QInterval, qint1: QInterval) → float
alkaid.trace.ops.bit_oprs.unary_bit_op(a: T, op: int, qint_from: QInterval, qint_to: QInterval) → T

alkaid.trace.ops.einsum_utils module

class alkaid.trace.ops.einsum_utils.EinsumRecipe

Bases: TypedDict

C: int
I: int
L0: int
L1: int
direct_sum_axis: tuple[tuple[int, ...], tuple[int, ...]]
in_transpose_idxs: tuple[tuple[int, ...], tuple[int, ...]]
out_interpert_shape: tuple[int, ...]
out_transpose_idxs: tuple[int, ...]
alkaid.trace.ops.einsum_utils.einsum(fn: str, input0: FVArray, input1: FVArray) → FVArray
alkaid.trace.ops.einsum_utils.einsum(fn: str, input0: FVArray, input1: ndarray[tuple[Any, ...], dtype[integer | floating]]) → FVArray
alkaid.trace.ops.einsum_utils.einsum(fn: str, input0: ndarray[tuple[Any, ...], dtype[integer | floating]], input1: FVArray) → FVArray
alkaid.trace.ops.einsum_utils.einsum(fn: str, input0: ndarray[tuple[Any, ...], dtype[integer | floating]], input1: ndarray[tuple[Any, ...], dtype[integer | floating]]) → ndarray[tuple[Any, ...], dtype[integer | floating]]
alkaid.trace.ops.einsum_utils.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:

einsum recipe; executed by _exec_einsum

Return type:

EinsumRecipe

alkaid.trace.ops.histogram module

Histogram and searchsorted operations for FVariable arrays.

alkaid.trace.ops.histogram.histogram(a: FVArray | NDArray, bins: int | NDArray = 10, range: tuple[float, float] | None = None, weights: FVArray | NDArray | None = None, density: Literal[False] = False) → tuple[FVArray | NDArray, NDArray]

Compute histogram using thermometer code counting.

Bin edges must be compile-time constants. Only internal edges are compared; out-of-range elements naturally land in the first/last bins.

Parameters:

  • a (FVArray or ndarray) – Input data. Flattened before processing.

  • bins (int or 1-D array-like) – If int: number of equal-width bins (range required). If array: monotonically increasing bin edges.

  • range ((float, float) or None) – Required when bins is an int.

  • weights (FVArray or ndarray or None) – Per-element weights (same shape as a). When given, each element contributes its weight instead of 1 to the bin count.

  • density (bool) – Not supported, raises ValueError if True.

Returns:

  • counts (FVArray or ndarray)

  • bin_edges (ndarray)

alkaid.trace.ops.histogram.searchsorted(a: FVArray | NDArray, v: FVArray | NDArray, side: str = 'left', sorter: str | NDArray | None = None) → FVArray | NDArray

Synthesisable searchsorted. Either a or v (or both) may be a FVArray.

Parameters:

  • a (FVArray or ndarray) – Sorted 1-D array of edges. May be symbolic (FVArray) or compile-time constants (plain ndarray).

  • v (FVArray or ndarray) – Values to search for.

  • side ({'left', 'right'}) – 'left': first index where a[i-1] < v <= a[i]. 'right': first index where a[i-1] <= v < a[i].

  • sorter (str or None) –

    Implementation strategy.

    • 'thermometer' — N parallel comparators + tree-sum. Works for both constant and symbolic edges. Good for short edge arrays where all comparisons fit in one pipeline stage.

    • 'bsearch' — binary search with msb_mux trees. Uses ceil(log2(N+1)) sequential stages; each stage maps to a pipeline stage. Works with both constant and symbolic edges.

    • None (default) — 'thermometer' when len(a) <= 8 or a is a FVArray, 'bsearch' otherwise.

alkaid.trace.ops.images module

Framework-neutral sliding-window patch extraction.

Pure numpy, works on both np.ndarray and alkaid.trace.FVArray. Used by both the keras and torch converter wrappers.

extract_patches and extract_patches_transposed share everything except the per-dim index formula and the output spatial size; the common tail is in _gather_patches.

alkaid.trace.ops.images.extract_patches(images: T, size: int | Sequence[int], strides: int | Sequence[int] | None = 1, dilation_rate: int | Sequence[int] = 1, padding: int | str | Sequence[int] | Sequence[tuple[int, int]] = 'valid', data_format: str = 'channels_last', pad_value: float = 0) → T

Forward sliding-window patch extraction.

For each output position o and kernel index k, gathers images[..., o*stride + k*dilation - pad_before, ...], masking positions that fall outside the original spatial range.

Parameters:

  • images – Array of shape (batch, *spatial, ch) for channels_last or (batch, ch, *spatial) for channels_first.

  • size – Per-spatial-dim kernel size. Scalar broadcasts to all spatial dims.

  • strides – Per-spatial-dim stride. None defaults to 1.

  • dilation_rate – Per-spatial-dim dilation.

  • padding – One of: 'valid', 'same' (keras ceil(in/stride) semantics), a single int (symmetric on every dim), a sequence of ints (symmetric per dim), or a sequence of (before, after) pairs.

  • data_format – 'channels_last' or 'channels_first'. Output is always channels-last form: (batch, *out_spatial, prod(size) * ch).

  • pad_value – Fill value for out-of-bounds positions.

alkaid.trace.ops.images.extract_patches_transposed(images: T, size: int | Sequence[int], strides: int | Sequence[int] = 1, dilation_rate: int | Sequence[int] = 1, padding: int | str | Sequence[int] | Sequence[tuple[int, int]] = 0, output_padding: int | Sequence[int] = 0, data_format: str = 'channels_last', pad_value: float = 0) → T

Gather-based patch extraction for transposed convolution.

Output spatial size per dim:

(in - 1) * stride - pad_before - pad_after + dilation * (size - 1) + output_padding + 1.

alkaid.trace.ops.quantization module

alkaid.trace.ops.quantization.quantize(x: T, k: ndarray[tuple[Any, ...], dtype[integer]] | integer | int, i: ndarray[tuple[Any, ...], dtype[integer]] | integer | int, f: ndarray[tuple[Any, ...], dtype[integer]] | integer | int, overflow_mode: str = 'WRAP', round_mode: str = 'TRN') → T
alkaid.trace.ops.quantization.relu(x: T, i: ndarray[tuple[Any, ...], dtype[integer]] | None = None, f: ndarray[tuple[Any, ...], dtype[integer]] | None = None, round_mode: str = 'TRN') → T

alkaid.trace.ops.reduce_utils module

class alkaid.trace.ops.reduce_utils.Packet(v)

Bases: object

alkaid.trace.ops.reduce_utils.argreduce(arr: FVArray, axis: int | Sequence[int] | None = None, keepdims: bool = False, minimize=True)

Reduction returning the index of the min or max element.

alkaid.trace.ops.reduce_utils.reduce(operator: Callable[[T, T], T], arr: TA, axis: int | Sequence[int] | None = None, keepdims: bool = False) → TA

Reduce the array by the operator over the specified axis.

alkaid.trace.ops.sorting module

alkaid.trace.ops.sorting.batcher_odd_even_merge_sort(a: ndarray[tuple[Any, ...], dtype[_ScalarT]], ascending: bool)

Sort an in-place power-of-two sequence with Batcher odd-even mergesort.

alkaid.trace.ops.sorting.cmp_swap(a: Sequence[FVariable] | NDArray, b: Sequence[FVariable] | NDArray, ascending: bool)
alkaid.trace.ops.sorting.sort(a: ndarray, axis: int | None = None, kind: str = 'batcher', aux_value: None = None) → FVArray
alkaid.trace.ops.sorting.sort(a: ndarray, axis: int | None = None, kind: str = 'batcher', aux_value: FVArray = None) → tuple[FVArray, FVArray]

Module contents