da4ml.trace.ops package

Submodules

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.

Module contents