Make the expectation functions traceable

[MPMPMPMPMPMPMP]

TLDR;

The idea of this PR is that we can maybe collect all the changes that are necessary to make the expectation functions traceable. Then we can jit a part of the __call__ method of the expectation function which would improve memory since jax also optimizes that under jit and it would allow us to checkpoint the computation of the expectation since it is very memory intensive to compute the RDM. I know there is also the checkpoint_ncon but this makes it more flexible.

the problematic part was this which is included in all expectation functions:

real_result = all(jnp.allclose(g, g.T.conj()) for g in gates)

This construct performs a Python-level loop and uses JAX arrays in a dynamic control context, which prevents JAX from tracing or staging out the function properly.
Despite multiple workarounds, none of them integrated cleanly with JAX’s tracing model or yielded good memory behavior.
My proposal would be to remove all of these and let the user handle the dtype in the model they define. Since the result array is also only not big ~ num gates there should be no problem with the memory

A example from my code

    @staticmethod
    @partial(jax.jit, static_argnames=("only_unique", "is_real"))
    def _accumulate_over_unitcell(
        peps_tensors,
        unitcell,
        working_up_gates,
        working_down_gates,
        *,
        only_unique: bool,
        is_real: bool,
    ):
        # initialize accumulator as a tuple to keep it purely functional
        zero = jnp.array(0.0) if is_real else jnp.array(0.0 + 0.0j)
        result = tuple(zero for _ in range(len(working_up_gates)))

        for x, iter_rows in unitcell.iter_all_rows(only_unique=only_unique):
            for y, view in iter_rows:
                # Get all 4 tensors in the 2x2 view
                tensors_i = view.get_indices((slice(0, 2, None), slice(0, 2, None)))
                tensors = [peps_tensors[i] for j in tensors_i for i in j]
                tensor_objs = [t for tl in view[:2, :2] for t in tl]

                # remat/checkpoint for memory; gates treated static via static_argnums
                step_result_down = jax.checkpoint(
                    calc_three_sites_triangle_without_bottom_left_multiple_gates
                )(tensors, tensor_objs, working_down_gates)
                step_result_up = jax.checkpoint(
                    calc_three_sites_triangle_without_top_right_multiple_gates
                )(tensors, tensor_objs, working_up_gates)

                # functional accumulation; ensure real dtype if requested
                incr = tuple(
                    ((sd.real if is_real else sd) + (su.real if is_real else su))
                    for sd, su in zip(step_result_down, step_result_up, strict=True)
                )
                result = tuple(r + a for r, a in zip(result, incr))
        return result

    def __call__(
        self,
        peps_tensors: Sequence[jnp.ndarray],
        unitcell: PEPS_Unit_Cell,
        spiral_vectors: Optional[Union[jnp.ndarray, Sequence[jnp.ndarray]]] = None,
        *,
        normalize_by_size: bool = True,
        only_unique: bool = True,
        return_single_gate_results: bool = False,
    ) -> Union[jnp.ndarray, List[jnp.ndarray]]:
        if self.is_spiral_peps:
            if spiral_vectors is None:
                raise ValueError(
                    "When using spiral iPEPS, spiral_vectors must be provided."
                )
            # if not isinstance(spiral_vectors, collections.abc.Sequence):
            #     spiral_vectors = (spiral_vectors,) * 3
            

            #[top-left, top-right, bottom-right]
            working_down_gates = tuple(
                apply_unitary(
                    h,
                    tuple(jnp.array(ri) for ri in ((0,0), (0,1), (1,1))),
                    spiral_vectors,
                    self._spiral_D,
                    self._spiral_sigma,
                    self.real_d,
                    3,
                    (0, 1, 2),
                    varipeps_config.spiral_wavevector_type,
                )
                for h in self.up_gates
            )
            #[top-left, bottom-left, bottom-right]
            working_up_gates = tuple(
                apply_unitary(
                    h,
                    tuple(jnp.array(ri) for ri in ((0,0), (1,0), (1,1))),
                    spiral_vectors,
                    self._spiral_D,
                    self._spiral_sigma,
                    self.real_d,
                    3,
                    (0, 1, 2),
                    varipeps_config.spiral_wavevector_type,
                )
                for h in self.down_gates
            )
        else:
            working_up_gates = self.up_gates
            working_down_gates = self.down_gates

        # Use jitted static method to perform the accumulation over the unitcell
        result = self._accumulate_over_unitcell(
            peps_tensors,
            unitcell,
            working_up_gates,
            working_down_gates,
            only_unique=only_unique,
            is_real=(self._result_type == jnp.float64),
        )

        if normalize_by_size:
            size = unitcell.get_len_unique_tensors() if only_unique else (unitcell.get_size()[0] * unitcell.get_size()[1])
            size = size * self.normalization_factor
            result = [r / size for r in result]

        if len(result) == 1:
            return result[0]
        else:
            return result

AI SUMMARY

This pull request simplifies the three_sites.py module by removing the real_result logic from the _three_site_triangle_workhorse function and its associated callers. This streamlines the computation of expectation values for three-site triangles, ensuring consistent output regardless of whether the gates are Hermitian.

Refactoring and code simplification:

• Removed the real_result argument from the _three_site_triangle_workhorse function signature and its usage in all calling functions, eliminating conditional logic based on gate Hermiticity. [1] [2]
• Deleted the calculation of the real_result variable (which checked for Hermitian gates) from all relevant calc_three_sites_triangle_without_*_multiple_gates functions. [1] [2] [3] [4]
• Updated all calls to _three_site_triangle_workhorse to remove the real_result parameter, further simplifying the function interfaces. [1] [2] [3] [4]

[JanLuca]

Yeah, that's right that this is a problem why we cannot trace over the expectation functions but why not just using jax.lax.cond? I can take a look into it but this should be manageable that one just use that to make this traceable

[MPMPMPMPMPMPMP]

I also tried that, but true_func and false_func need to have the same dtype. Which is obviously not the case.

The true_computation must take in a single argument of type and will be invoked with true_operand which must be of the same type. The false_computation must take in a single argument of type and will be invoked with false_operand which must be of the same type. The type of the returned value of true_computation and false_computation must be the same.

taken from https://openxla.org/xla/operation_semantics#conditional this is wrapped in jax.lax.cond