pdft.circuit.optimize_code_cached

pdft.circuit.optimize_code_cached#

pdft.circuit.optimize_code_cached(subscripts, *shapes)[source]#

Return a jit-compiled einsum closure bound to fixed subscripts and shapes.

On first call for a given (subscripts, shapes) key, computes an optimal contraction path via jnp.einsum_path on dummy arrays of the supplied shapes, then returns a jit-compiled closure. Subsequent calls return the identical cached closure.

Parameters:
  • subscripts (str) – einsum subscripts string, e.g., “ijk,jlk->ilk”.

  • *shapes (tuple[int, ...]) – Shape of each operand, in order matching the subscripts.

Returns:

Callable accepting the operands and returning the contraction result.

Return type:

Callable[[…], Array]