module testing.einsum.einsum_impl

Short summary

module mlprodict.testing.einsum.einsum_impl

Main functions decomposing einsum computation into more simple functions.

source on GitHub

Functions

function	truncated documentation
_apply_einsum_matmul	Decomposes the generic matrix multiplication into numpy operations depending on the operator to use for matrix multiplication …
_apply_squeeze_transpose	Puts output dimension in the expected order.
_apply_transpose_reshape	Put all dimensions in the same order.
_basic_verification
_decompose_einsum_equation_simple	Applies strategy simple, numpy defined in by function decompose_einsum_equation().
analyse_einsum_equation	Analyses an einsum equation.
apply_einsum_sequence	Applies a sequence of operations on a list of inputs. The sequence of operations is produced by function decompose_einsum_equation(). …
decompose_einsum_equation	Decomposes an equation used in numpy.einsum knowing the input shapes. It returns a sequence of operations …
is_transpose_identity	Tells if the permutation perm does nothing (itentity).

Documentation

Main functions decomposing einsum computation into more simple functions.

source on GitHub

mlprodict.testing.einsum.einsum_impl._apply_einsum_matmul(fd, op1, op2, axes, left, right, ndim, op_matmul, row1, row2, verbose=False)

Decomposes the generic matrix multiplication into numpy operations depending on the operator to use for matrix multiplication op_matmul (see decompose_einsum_equation).

source on GitHub

mlprodict.testing.einsum.einsum_impl._apply_squeeze_transpose(op, row_last, row_output)

Puts output dimension in the expected order.

source on GitHub

mlprodict.testing.einsum.einsum_impl._apply_transpose_reshape(op, row)

Put all dimensions in the same order.

Parameters:

• op – integer (for one input) or an operator

• row – letter involved in this input (as a vector of binaries)

Returns:

last created operator

source on GitHub

mlprodict.testing.einsum.einsum_impl._basic_verification(lengths, shapes, equation)

mlprodict.testing.einsum.einsum_impl._decompose_einsum_equation_simple(equation, *shapes, verbose=False, op_matmul='matmul')

Applies strategy simple, numpy defined in by function decompose_einsum_equation.

Parameters:

op_matmul – which operator to use for matrix multiplication, a single operator matmul, or batch_dot with transposes, reduce_sum, or just dot

source on GitHub

mlprodict.testing.einsum.einsum_impl.analyse_einsum_equation(equation)

Analyses an einsum equation.

Parameters:

equation – numpy.einsum equation

Returns:

three results, list of letters, a matrix (see below), lengths of each components, duplicates

The returned a matrix is defined as follows:

source on GitHub

mlprodict.testing.einsum.einsum_impl.apply_einsum_sequence(seq, *inputs, verbose=False, **kwargs)

Applies a sequence of operations on a list of inputs. The sequence of operations is produced by function decompose_einsum_equation.

Parameters:

• seq – sequence of operations

• inputs – inputs

• kwargs – additional parameters, see apply_sequence.

Returns:

output

<<<

import numpy
from mlprodict.testing.einsum import (
    decompose_einsum_equation, apply_einsum_sequence)

m1 = numpy.arange(2 * 2 * 2).reshape((2, 2, 2)) + 10
m2 = numpy.arange(4).reshape((2, 2)) + 100
m3 = numpy.arange(8).reshape((2, 2, 2)) + 1000

seq = decompose_einsum_equation("bac,cd,def->ebc")
res = apply_einsum_sequence(seq, m1, m2, m3)
print(res)

>>>

    [[[ 8866198  9864696]
      [12090270 13152928]]
    
     [[ 8883886  9884376]
      [12114390 13179168]]]

See notebook Einsum decomposition.

source on GitHub

mlprodict.testing.einsum.einsum_impl.decompose_einsum_equation(equation, *shapes, strategy='simple', clean=False, verbose=False)

Decomposes an equation used in numpy.einsum knowing the input shapes. It returns a sequence of operations to do to compute the results.

Parameters:

• equation – a string

• shapes – sequence of input shapes

• strategy – there are different way to decompose the equation, this parameters defines the way to do it (see below)

• clean – clean the unnecessary node in the graph

• verbose – verbosity

Returns:

instance of GraphEinsumSubOp

About strategy:

• ‘simple’: align all dimensions in the alphabetical order, some generic matrix multiplication remains implemented with numpy.einsum but only with two matrices aligned on the same dimension (see numpy_extended_dot)

• ‘numpy’: same as simple but the decomposition does not use numpy.einsum anymore but only multiplication or matrix multiplication merged into a single operator called batch_dot (see numpy_extended_dot_matrix)

Available operations: expand_dims, transpose, matmul, reduce_sum, id, squeeze, diagonal. It analyses an equation and produces a graph where node are instance of class EinsumSubOp.

<<<

from mlprodict.testing.einsum import decompose_einsum_equation
seq = decompose_einsum_equation("bac,cd,def->ebc")
for op in seq:
    print(op)

>>>

    EinsumSubOp('id', 0, )
    EinsumSubOp('expand_dims', EinsumSubOp('id', 0, ), axes=((3, 3), (3, 4), (3, 5)))
    EinsumSubOp('transpose', EinsumSubOp('expand_dims', EinsumSubOp('id', 0, ), axes=((3, 3), (3, 4), (3, 5))), perm=(1, 0, 2, 3, 4, 5))
    EinsumSubOp('reduce_sum', EinsumSubOp('transpose', EinsumSubOp('expand_dims', EinsumSubOp('id', 0, ), axes=((3, 3), (3, 4), (3, 5))), perm=(1, 0, 2, 3, 4, 5)), axes=(0,))
    EinsumSubOp('id', 1, )
    EinsumSubOp('expand_dims', EinsumSubOp('id', 1, ), axes=((0, 0), (0, 1), (2, 4), (2, 5)))
    EinsumSubOp('matmul', EinsumSubOp('reduce_sum', EinsumSubOp('transpose', EinsumSubOp('expand_dims', EinsumSubOp('id', 0, ), axes=((3, 3), (3, 4), (3, 5))), perm=(1, 0, 2, 3, 4, 5)), axes=(0,)), EinsumSubOp('expand_dims', EinsumSubOp('id', 1, ), axes=((0, 0), (0, 1), (2, 4), (2, 5))), axes=(), left=(1, 2), right=(2, 3), ndim=6)
    EinsumSubOp('id', 2, )
    EinsumSubOp('expand_dims', EinsumSubOp('id', 2, ), axes=((0, 0), (0, 1), (0, 2)))
    EinsumSubOp('reduce_sum', EinsumSubOp('expand_dims', EinsumSubOp('id', 2, ), axes=((0, 0), (0, 1), (0, 2))), axes=(5,))
    EinsumSubOp('matmul', EinsumSubOp('matmul', EinsumSubOp('reduce_sum', EinsumSubOp('transpose', EinsumSubOp('expand_dims', EinsumSubOp('id', 0, ), axes=((3, 3), (3, 4), (3, 5))), perm=(1, 0, 2, 3, 4, 5)), axes=(0,)), EinsumSubOp('expand_dims', EinsumSubOp('id', 1, ), axes=((0, 0), (0, 1), (2, 4), (2, 5))), axes=(), left=(1, 2), right=(2, 3), ndim=6), EinsumSubOp('reduce_sum', EinsumSubOp('expand_dims', EinsumSubOp('id', 2, ), axes=((0, 0), (0, 1), (0, 2))), axes=(5,)), axes=(3,), left=(1, 2), right=(4,), ndim=6)
    EinsumSubOp('transpose', EinsumSubOp('matmul', EinsumSubOp('matmul', EinsumSubOp('reduce_sum', EinsumSubOp('transpose', EinsumSubOp('expand_dims', EinsumSubOp('id', 0, ), axes=((3, 3), (3, 4), (3, 5))), perm=(1, 0, 2, 3, 4, 5)), axes=(0,)), EinsumSubOp('expand_dims', EinsumSubOp('id', 1, ), axes=((0, 0), (0, 1), (2, 4), (2, 5))), axes=(), left=(1, 2), right=(2, 3), ndim=6), EinsumSubOp('reduce_sum', EinsumSubOp('expand_dims', EinsumSubOp('id', 2, ), axes=((0, 0), (0, 1), (0, 2))), axes=(5,)), axes=(3,), left=(1, 2), right=(4,), ndim=6), perm=(0, 4, 1, 3, 2, 5))
    EinsumSubOp('squeeze', EinsumSubOp('transpose', EinsumSubOp('matmul', EinsumSubOp('matmul', EinsumSubOp('reduce_sum', EinsumSubOp('transpose', EinsumSubOp('expand_dims', EinsumSubOp('id', 0, ), axes=((3, 3), (3, 4), (3, 5))), perm=(1, 0, 2, 3, 4, 5)), axes=(0,)), EinsumSubOp('expand_dims', EinsumSubOp('id', 1, ), axes=((0, 0), (0, 1), (2, 4), (2, 5))), axes=(), left=(1, 2), right=(2, 3), ndim=6), EinsumSubOp('reduce_sum', EinsumSubOp('expand_dims', EinsumSubOp('id', 2, ), axes=((0, 0), (0, 1), (0, 2))), axes=(5,)), axes=(3,), left=(1, 2), right=(4,), ndim=6), perm=(0, 4, 1, 3, 2, 5)), axes=(0, 3, 5))

It can be better displayed as the following.

See notebook Einsum decomposition.

source on GitHub

mlprodict.testing.einsum.einsum_impl.is_transpose_identity(perm)

Tells if the permutation perm does nothing (itentity).

Parameters:

perm – permutation

Returns:

boolean

source on GitHub