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

    @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