yobx.helpers._einsum.einsum_impl_classes

class yobx.helpers._einsum.einsum_impl_classes.EinsumSubOp(full_dim: int, name: str, *inputs: int | EinsumSubOp, **kwargs: Any)

Defines a sub operation used in Einsum decomposition.

Parameters:

• full_dim – dimension of the result

• name – name (reshape, transpose, reduce_sum, matmul, id, squeeze, diagonal, mul, batch_dot)

• inputs – inputs

• kwargs – arguments

Operator suffixed by _mm (transpose_mm, reduce_sum_mm) are equivalent to the same operator without the suffix but takes two inputs and only changes the first one.

Attributes _info summarizes the known information about dimensions. Many of them are empty because inserted. Value 1 means it was the case, 2 means it is a plain dimension.

add_info(**kwargs: Any)

Adds information to the node.

Parameters:

kwargs – dictionary

apply(data: Dict[int, Any], verbose: bool = False, **kwargs: Any) → Any

Applies one operator on the data.

Parameters:

• data – dictionary storing the results

• verbose – prints out intermediate results

• kwargs – additional parameters, see methods _apply*

Returns:

output

Known additional parameters:

• ‘matmul_impl’: if None calls numpy.einsum() through numpy_extended_dot (wrong) (default) or ‘py’ to call numpy_extended_dot_python instead.

compute_output_row(row: ndarray, row2: ndarray | None = None, ab: bool = False, verbose: bool = False)

Updates row based on the operator.

dot_label() → str | None

Displays some information useful to understand the operator.

get_dot_kind() → str

Every matrix multiplication can be either:

• a simple multiplication (M) (undetected)

• a 2D matrix multiplication (11)

• a broadcasted matrix multiplication (N1 or 1N)

• a batch matrix multiplication (NN)

This method returns which kind it is.

to_onnx(names: Dict[int, str], opset: int | None, verbose: bool = False, **kwargs: Any) → Iterable[NodeProto | TensorProto]

Converts this node into ONNX. Enumerates all ONNX node which participate to the conversion. The last one is the final output.

Parameters:

• names – dictionary where to find already converted name

• opset – opset

• verbose – prints out intermediate results

• kwargs – additional parameter for the conversion

Returns:

output

class yobx.helpers._einsum.einsum_impl_classes.GraphEinsumSubOp(letters: str, mat: ndarray, lengths: List[int], duplicates: List[Dict[str, Any] | None])

Class gathering all nodes produced to explicit einsum operators.

Parameters:

• letters – list of distinct letters

• mat – matrix, see analyse_einsum_equation

• lengths – lengths of every input

• duplicates – see analyse_einsum_equation

append(op: int | EinsumSubOp) → EinsumSubOp | None

Adds one input or result.

Parameters:

op – integer (an input) or an instance of EinsumSubOp.

Returns:

op or None if op is an integer

apply_sequence(*inputs: Any, verbose: bool = False, **kwargs: Any) → Any

Applies a sequence of operations on a list of inputs.

Parameters:

• inputs – inputs

• verbose – prints out intermediate results

• kwargs – additional parameters, see apply.

Returns:

output

clean_unused_nodes(verbose: bool = False)

Cleans nodes with unused outputs.

Parameters:

verbose – display intermediate information

mark(i: int, op: int | EinsumSubOp)

Marks one input or result as an intermediate result after a full einsum step.

Parameters:

• i – a position

• op – an instance of EinsumSubOp.

mark_last_node()

Marks the last node as the final output.

remove_duplicate_transpose(verbose: bool = False)

Removes consecutive transpose by merging them.

Parameters:

verbose – display intermediate information

simplify_mm_nodes(verbose: bool = False)

Node name suffixed by mm are an artifact to keep the graph consistent while building it. They can now be replaced by the equivalent node without suffix mm.

Parameters:

verbose – display intermediate information

to_dot(**kwargs: Any) → str

Produces a graph in dot.

Parameters:

kwargs – additional graph option

Returns:

string

to_onnx(output: str, *inputs: Any, dtype: Any | None = None, verbose: bool = False, opset: int | None = None, **kwargs: Any) → ModelProto

Converts the graph into ONNX.

Parameters:

• output – output name

• inputs – input names

• dtype – type used for all operators

• opset – desired opset, None for the last one

• verbose – display intermediate operators

• kwargs – additional parameter to use when building the ONNX graph, list of supported parameters: name, ir_version, producer_name, producer_version, initializer

Returns:

ONNX graph

Not all graphs can be converted into ONNX. Only graphs produced with strategy=’numpy’ can be converted otherwise the following error shows up:

NotImplementedError: to_onnx not implemented for 'matmul'.

yobx.helpers._einsum.einsum_impl_classes.single_axes(axes: Any) → List[int] | None

axes contains positive values, then it is the position of this axis in the original matrix, otherwise it is -1 meaning this axis is an added single dimension to align all the dimensions based on the einsum equation.

Parameters:

axes – axes described above

Returns:

list of integer in set {1, 2}, 1 for a single axis, 2 otherwise