da4ml.trace.ops package

Submodules

da4ml.trace.ops.bit_oprs module

da4ml.trace.ops.bit_oprs.binary_bit_op(a: FixedVariable, b: FixedVariable, op: int, *args, **kwargs) → FixedVariable
da4ml.trace.ops.bit_oprs.binary_bit_op(a: float, b: float, op: int, qint0: QInterval, qint1: QInterval) → float

da4ml.trace.ops.bit_oprs.unary_bit_op(a: T, op: int, qint_from: QInterval, qint_to: QInterval) → T

da4ml.trace.ops.einsum_utils module

class da4ml.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, ...]

da4ml.trace.ops.einsum_utils.einsum(fn: str, input0: FixedVariableArray, input1: FixedVariableArray) → FixedVariableArray
da4ml.trace.ops.einsum_utils.einsum(fn: str, input0: FixedVariableArray, input1: ndarray[tuple[Any, ...], dtype[integer | floating]]) → FixedVariableArray
da4ml.trace.ops.einsum_utils.einsum(fn: str, input0: ndarray[tuple[Any, ...], dtype[integer | floating]], input1: FixedVariableArray) → FixedVariableArray
da4ml.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]]

da4ml.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

da4ml.trace.ops.quantization module

da4ml.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

da4ml.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

da4ml.trace.ops.reduce_utils module

class da4ml.trace.ops.reduce_utils.Packet(v): Bases: object

da4ml.trace.ops.reduce_utils.reduce(operator: Callable[[T, T], T], x: TA, axis: int | Sequence[int] | None = None, keepdims: bool = False) → TA: Reduce the array by summing over the specified axis.

da4ml.trace.ops.sorting module

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

##copy-paste from wikipedia https://en.wikipedia.org/wiki/Batcher_odd%E2%80%93even_mergesort; apparently it works

note: the input sequence is indexed from 0 to (n-1) for p = 1, 2, 4, 8, … # as long as p < n

for k = p, p/2, p/4, p/8, … # as long as k >= 1

for j = mod(k,p) to (n-1-k) with a step size of 2k

for i = 0 to min(k-1, n-j-k-1) with a step size of 1

if floor((i+j) / (p*2)) == floor((i+j+k) / (p*2))
compare and sort elements (i+j) and (i+j+k)

da4ml.trace.ops.sorting.cmp_swap(a: Sequence[FixedVariable] | NDArray, b: Sequence[FixedVariable] | NDArray, ascending: bool)

da4ml.trace.ops.sorting.sort(a: ndarray, axis: int | None = None, kind: str = 'batcher', aux_value: None = None) → ndarray
da4ml.trace.ops.sorting.sort(a: FixedVariableArray, axis: int | None = None, kind: str = 'batcher', aux_value: None = None) → FixedVariableArray
da4ml.trace.ops.sorting.sort(a: FixedVariableArray, axis: int | None = None, kind: str = 'batcher', aux_value: FixedVariableArray = None) → tuple[FixedVariableArray, FixedVariableArray]