:py:mod:`cotengra.contract` =========================== .. py:module:: cotengra.contract .. autoapi-nested-parse:: Functionality relating to actually contracting. Module Contents --------------- Classes ~~~~~~~ .. autoapisummary:: cotengra.contract.Contractor cotengra.contract.CuQuantumContractor Functions ~~~~~~~~~ .. autoapisummary:: cotengra.contract.set_default_implementation cotengra.contract.get_default_implementation cotengra.contract.default_implementation cotengra.contract._sanitize_equation cotengra.contract._parse_einsum_single cotengra.contract._parse_eq_to_pure_multiplication cotengra.contract._parse_eq_to_batch_matmul cotengra.contract._einsum_single cotengra.contract._do_contraction_via_bmm cotengra.contract.einsum cotengra.contract.gen_nice_inds cotengra.contract._parse_tensordot_axes_to_matmul cotengra.contract.tensordot cotengra.contract.extract_contractions cotengra.contract.make_contractor Attributes ~~~~~~~~~~ .. autoapisummary:: cotengra.contract.DEFAULT_IMPLEMENTATION .. py:data:: DEFAULT_IMPLEMENTATION :value: 'auto' .. py:function:: set_default_implementation(impl) .. py:function:: get_default_implementation() .. py:function:: default_implementation(impl) Context manager for temporarily setting the default implementation. .. py:function:: _sanitize_equation(eq) Get the input and output indices of an equation, computing the output implicitly as the sorted sequence of every index that appears exactly once if it is not provided. .. py:function:: _parse_einsum_single(eq, shape) Cached parsing of a single term einsum equation into the necessary sequence of arguments for axes diagonals, sums, and transposes. .. py:function:: _parse_eq_to_pure_multiplication(a_term, shape_a, b_term, shape_b, out) If there are no contracted indices, then we can directly transpose and insert singleton dimensions into ``a`` and ``b`` such that (broadcast) elementwise multiplication performs the einsum. No need to cache this as it is within the cached ``_parse_eq_to_batch_matmul``. .. py:function:: _parse_eq_to_batch_matmul(eq, shape_a, shape_b) Cached parsing of a two term einsum equation into the necessary sequence of arguments for contracttion via batched matrix multiplication. The steps we need to specify are: 1. Remove repeated and trivial indices from the left and right terms, and transpose them, done as a single einsum. 2. Fuse the remaining indices so we have two 3D tensors. 3. Perform the batched matrix multiplication. 4. Unfuse the output to get the desired final index order. .. py:function:: _einsum_single(eq, x, backend=None) Einsum on a single tensor, via three steps: diagonal selection (via advanced indexing), axes summations, transposition. The logic for each is cached based on the equation and array shape, and each step is only performed if necessary. .. py:function:: _do_contraction_via_bmm(a, b, eq_a, eq_b, new_shape_a, new_shape_b, new_shape_ab, perm_ab, pure_multiplication, backend) .. py:function:: einsum(eq, a, b=None, *, backend=None) Perform arbitrary single and pairwise einsums using only `matmul`, `transpose`, `reshape` and `sum`. The logic for each is cached based on the equation and array shape, and each step is only performed if necessary. :param eq: The einsum equation. :type eq: str :param a: The first array to contract. :type a: array_like :param b: The second array to contract. :type b: array_like, optional :param backend: The backend to use for array operations. If ``None``, dispatch automatically based on ``a`` and ``b``. :type backend: str, optional :rtype: array_like .. py:function:: gen_nice_inds() Generate the indices from [a-z, A-Z, reasonable unicode...]. .. py:function:: _parse_tensordot_axes_to_matmul(axes, shape_a, shape_b) Parse a tensordot specification into the necessary sequence of arguments for contracttion via matrix multiplication. This just converts ``axes`` into an ``einsum`` eq string then calls ``_parse_eq_to_batch_matmul``. .. py:function:: tensordot(a, b, axes=2, *, backend=None) Perform a tensordot using only `matmul`, `transpose`, `reshape`. The logic for each is cached based on the equation and array shape, and each step is only performed if necessary. :param a: The arrays to contract. :type a: array_like :param b: The arrays to contract. :type b: array_like :param axes: The number of axes to contract, or the axes to contract. If an int, the last ``axes`` axes of ``a`` and the first ``axes`` axes of ``b`` are contracted. If a tuple, the axes to contract for ``a`` and ``b`` respectively. :type axes: int or tuple of (sequence[int], sequence[int]) :param backend: The backend to use for array operations. If ``None``, dispatch automatically based on ``a`` and ``b``. :type backend: str or None, optional :rtype: array_like .. py:function:: extract_contractions(tree, order=None, prefer_einsum=False) Extract just the information needed to perform the contraction. :param order: Supplied to :meth:`ContractionTree.traverse`. :type order: str or callable, optional :param prefer_einsum: Prefer to use ``einsum`` for pairwise contractions, even if ``tensordot`` can perform the contraction. :type prefer_einsum: bool, optional :returns: **contractions** -- A tuple of tuples, each containing the information needed to perform a pairwise contraction. Each tuple contains: - ``p``: the parent node, - ``l``: the left child node, - ``r``: the right child node, - ``tdot``: whether to use ``tensordot`` or ``einsum``, - ``arg``: the argument to pass to ``tensordot`` or ``einsum`` i.e. ``axes`` or ``eq``, - ``perm``: the permutation required after the contraction, if any (only applies to tensordot). If both ``l`` and ``r`` are ``None``, the the operation is a single term simplification performed with ``einsum``. :rtype: tuple .. py:class:: Contractor(contractions, strip_exponent=False, check_zero=False, implementation='auto', backend=None, progbar=False) Default cotengra network contractor. :param contractions: The sequence of contractions to perform. Each contraction should be a tuple containing: - ``p``: the parent node, - ``l``: the left child node, - ``r``: the right child node, - ``tdot``: whether to use ``tensordot`` or ``einsum``, - ``arg``: the argument to pass to ``tensordot`` or ``einsum`` i.e. ``axes`` or ``eq``, - ``perm``: the permutation required after the contraction, if any (only applies to tensordot). e.g. built by calling ``extract_contractions(tree)``. :type contractions: tuple[tuple] :param strip_exponent: If ``True``, eagerly strip the exponent (in log10) from intermediate tensors to control numerical problems from leaving the range of the datatype. This method then returns the scaled 'mantissa' output array and the exponent separately. :type strip_exponent: bool, optional :param check_zero: If ``True``, when ``strip_exponent=True``, explicitly check for zero-valued intermediates that would otherwise produce ``nan``, instead terminating early if encounteredand returning ``(0.0, 0.0)``. :type check_zero: bool, optional :param backend: What library to use for ``tensordot``, ``einsum`` and ``transpose``, it will be automatically inferred from the input arrays if not given. :type backend: str, optional :param progbar: Whether to show a progress bar. :type progbar: bool, optional .. py:attribute:: __slots__ :value: ('contractions', 'strip_exponent', 'check_zero', 'implementation', 'backend', 'progbar', '__weakref__') .. py:method:: __call__(*arrays, **kwargs) Contract ``arrays`` using operations listed in ``contractions``. :param arrays: The arrays to contract. :type arrays: sequence of array-like :param kwargs: Override the default settings for this contraction only. :type kwargs: dict :returns: * **output** (*array*) -- The contracted output, it will be scaled if ``strip_exponent==True``. * **exponent** (*float*) -- The exponent of the output in base 10, returned only if ``strip_exponent==True``. .. py:class:: CuQuantumContractor(tree, handle_slicing=False, autotune=False, **kwargs) .. py:method:: setup(*arrays) .. py:method:: __call__(*arrays, check_zero=False, backend=None, progbar=False) .. py:method:: __del__() .. py:function:: make_contractor(tree, order=None, prefer_einsum=False, strip_exponent=False, check_zero=False, implementation=None, autojit=False, progbar=False) Get a reusable function which performs the contraction corresponding to ``tree``. The various options provide defaults that can also be overrode when calling the standard contractor. :param tree: The contraction tree. :type tree: ContractionTree :param order: Supplied to :meth:`ContractionTree.traverse`, the order in which to perform the pairwise contractions given by the tree. :type order: str or callable, optional :param prefer_einsum: Prefer to use ``einsum`` for pairwise contractions, even if ``tensordot`` can perform the contraction. :type prefer_einsum: bool, optional :param strip_exponent: If ``True``, the function will strip the exponent from the output array and return it separately. :type strip_exponent: bool, optional :param check_zero: If ``True``, when ``strip_exponent=True``, explicitly check for zero-valued intermediates that would otherwise produce ``nan``, instead terminating early if encountered and returning ``(0.0, 0.0)``. :type check_zero: bool, optional :param implementation: What library to use to actually perform the contractions. Options are - "auto": let cotengra choose - "autoray": dispatch with autoray, using the ``tensordot`` and ``einsum`` implementation of the backend - "cotengra": use the ``tensordot`` and ``einsum`` implementation of cotengra, which is based on batch matrix multiplication. This is faster for some backends like numpy, and also enables libraries which don't yet provide ``tensordot`` and ``einsum`` to be used. - "cuquantum": use the cuquantum library to perform the whole contraction (not just individual contractions). - tuple[callable, callable]: manually supply the ``tensordot`` and ``einsum`` implementations to use. :type implementation: str or tuple[callable, callable], optional :param autojit: If ``True``, use :func:`autoray.autojit` to compile the contraction function. :type autojit: bool, optional :param progbar: Whether to show progress through the contraction by default. :type progbar: bool, optional :returns: **fn** -- The contraction function, with signature ``fn(*arrays)``. :rtype: callable