ICOjax.py

# %%
import jax 
import jax.numpy as jnp
from jax import lax
import numpy as np
from functools import partial
from tqdm import tqdm
import optax
from jaxopt import BFGS
import time
from jax import random
import secrets
import os
import dumpy as dp

from tensors import jax_kron, permuteHilbert, partial_trace, partial_trace_with_identity, partial_trace_with_identity_numpy
from utils import find_permutation, chop
from quantumops import Lv_projector_process_2parties, Lv_projector_process_3parties, Lv_projector_channel
from jaxparams import param01_to_hermitian, make_operator_psd, hermitian_to_param01, give_Choi



@partial(jax.jit, static_argnums=(1, 2, 3, 4, 7))
def give_rank1_Choi(theta, dim_in, dim_out, nof_outcomes, nof_settings, ini_rank1_Choi, unitary_generators, nof_generators):
    dim_out_with_outcomes = dim_out * nof_outcomes
    dim = dim_in * dim_out_with_outcomes
    params_each_choi = nof_generators
    # trace_Choi_x_value = dim_in

    # Initialize penalty as a JAX scalar
    Choi_list = []
    penalty = jnp.float32(0.0)

    # Loop over settings
    for x in range(nof_settings):
        start_idx = x * params_each_choi
        end_idx = (x + 1) * params_each_choi

        # Extract parameters for the current Choi
        current_params = theta[start_idx:end_idx]

        # Convert parameters to a Hermitian matrix
        U = jnp.scipy.linalg.expm(-1j * jnp.dot(unitary_generators, current_params).reshape((dim, dim)))
        Choi = U @ ini_rank1_Choi @ U.conj().T

        # Project onto a valid channel space
        # Choi = Lv_projector_channel(Choi, dim_in, dim_out_with_outcomes)

        # Normalize Choi so that its trace = trace_Choi_x_value
        # Choi_trace = jnp.trace(Choi)
        # Choi = (trace_Choi_x_value / Choi_trace) * Choi

        # Update penalty: penalty += logdet(Choi)
        # jnp.linalg.logdet returns a real scalar if Choi is Hermitian positive-definite.
        
        # sign, logdet = jnp.linalg.slogdet(Choi)
        # penalty += logdet
        # Choi = make_operator_psd(Choi, trace_Choi_x_value, dim)

        # Collect Choi
        Choi_list.append(Choi)

    # Stack all Choi matrices along a new axis (e.g. axis=0)
    Choi_x = jnp.stack(Choi_list, axis=0)

    return Choi_x, penalty

@partial(jax.jit, static_argnums=(1,2,3))
def extract_Choi_xa(Choi_x_single, dim_in, dim_out, nof_outcomes):
    """
    Extracts the Choi_xa blocks for a single Choi_x.
    
    Args:
        Choi_x_single: A single Choi matrix of shape (dim_in*dim_out*nof_outcomes, dim_in*dim_out*nof_outcomes).
        dim_in: Input dimension
        dim_out: Output dimension
        nof_outcomes: Number of outcomes
    
    Returns:
        A JAX array of shape (nof_outcomes, dim_in*dim_out, dim_in*dim_out).
    """
    # Reshape to 6D: (dim_in, dim_out, nof_outcomes, dim_in, dim_out, nof_outcomes)
    reshaped = jnp.reshape(Choi_x_single, (dim_in, dim_out, nof_outcomes, dim_in, dim_out, nof_outcomes))

    # Transpose so that the two nof_outcomes axes are last, making it easier to take a "diagonal"
    # Reorder to: (dim_in, dim_out, dim_in, dim_out, nof_outcomes, nof_outcomes)
    # Original:    (0,      1,      2,      3,      4,             5)
    # Target:      (0,      1,      3,      4,      2,             5)
    transposed = jnp.transpose(reshaped, (0, 1, 3, 4, 2, 5))

    # Take the diagonal along the last two axes (both of length nof_outcomes)
    # This extracts (dim_in, dim_out, dim_in, dim_out, nof_outcomes) where the last dimension corresponds to 'a'
    diags = jnp.diagonal(transposed, axis1=4, axis2=5)

    # Move the nof_outcomes axis (currently last) to the front
    diags = jnp.moveaxis(diags, -1, 0)
    # Now shape is (nof_outcomes, dim_in, dim_out, dim_in, dim_out)

    # Finally, reshape each block to (dim_in*dim_out, dim_in*dim_out)
    return diags.reshape(nof_outcomes, dim_in*dim_out, dim_in*dim_out)

@partial(jax.jit, static_argnums=(1,2,3))
def extract_Choi_xa_all(Choi_x, dim_in, dim_out, nof_outcomes):
    """
    Extracts the Choi_xa blocks for all x in Choi_x.
    
    Args:
        Choi_x: A JAX array of shape (nof_settings, dim_in*dim_out*nof_outcomes, dim_in*dim_out*nof_outcomes)
                containing Choi matrices for multiple settings.
        dim_in, dim_out, nof_outcomes: As above.
        nof_settings: Number of settings.
    
    Returns:
        A JAX array of shape (nof_settings, nof_outcomes, dim_in*dim_out, dim_in*dim_out).
    """
    # We can use vmap to apply extract_Choi_xa over the first dimension (settings dimension) of Choi_x
    return jax.vmap(extract_Choi_xa, in_axes=(0, None, None, None))(Choi_x, dim_in, dim_out, nof_outcomes)


@partial(jax.jit, static_argnums=(1, 2))
def give_processmatrix(theta, dims_in, dims_out):
    dim_in = np.prod(dims_in)
    dim_out = np.prod(dims_out)
    dim = dim_in * dim_out
    params = dim**2
    trace_W_x_value = dim_in

    # Initialize penalty as a JAX scalar
    penalty = jnp.float32(0.0)

    # Extract parameters for the current Choi
    current_params = theta[:params]

    # Convert parameters to a Hermitian matrix
    W = param01_to_hermitian(current_params, dim)

    # Project onto a valid channel space
    _dims = []
    for _dim_in, _dim_out in zip(dims_in, dims_out):
        _dims += [_dim_in, _dim_out]
    W = Lv_projector_process_2parties(W, _dims[0], _dims[1], _dims[2], _dims[3])

    # Normalize Choi so that its trace = trace_Choi_x_value
    W_trace = jnp.trace(W)
    W = (trace_W_x_value / W_trace) * W


    # jnp.linalg.logdet returns a real scalar if Choi is Hermitian positive-definite.
    # sign, logdet = jnp.linalg.slogdet(W)
    # penalty += logdet
    W = make_operator_psd(W, trace_W_x_value, dim)

    return W, penalty

@partial(jax.jit, static_argnums=(1, 2, 3, 4, 5, 6, 7, 8))
def give_probability(theta, na, nb, nx, ny, dim_A_in, dim_A_out, dim_B_in, dim_B_out):
    nof_params_A = (dim_A_in * dim_A_out * na)**2 * nx
    nof_params_B = (dim_B_in * dim_B_out * nb)**2 * ny
    nof_params_W = (dim_A_in * dim_A_out * dim_B_in * dim_B_out)**2
    
    params_A = theta[:nof_params_A]
    params_B = theta[nof_params_A:nof_params_A+nof_params_B]
    params_W = theta[nof_params_A+nof_params_B:]
    
    Choi_x, penalty_A = give_Choi(params_A, dim_A_in, dim_A_out, na, nx)
    Choi_y, penalty_B = give_Choi(params_B, dim_B_in, dim_B_out, nb, ny)
    W, penalty_W = give_processmatrix(params_W, (dim_A_in, dim_A_out), (dim_B_in, dim_B_out))
    
    Choi_xa = extract_Choi_xa_all(Choi_x, dim_A_in, dim_A_out, na)
    Choi_yb = extract_Choi_xa_all(Choi_y, dim_B_in, dim_B_out, nb)
    
    # # Vectorize over outcomes and settings for party A
    # def compute_A_block(xa):
    #     # xa has shape (na, dimA, dimA)
    #     # vmap over a to produce a function that given a particular Choi_yb and W produces a full slice
    #     return jax.vmap(lambda A: jax.vmap(
    #         lambda YB: jax.vmap(lambda yb: jnp.real(jnp.trace(W @ jax_kron(A, yb))))(YB),
    #         in_axes=(None,)
    #     ))(Choi_yb)
    
    # prob_abxy = jax.vmap(compute_A_block, in_axes=(0,))(Choi_xa)
    # # prob_abxy now has shape (nx, na, nb, ny)
    # # reshape to (na, nb, nx, ny)
    # prob_abxy = jnp.transpose(prob_abxy, (1, 2, 0, 3))
    
    # Create a probability array
    # shape is (na, nb, nx, ny)
    # prob_abxy = jnp.zeros((na, nb, nx, ny), dtype=jnp.float32)
    # for a, b in np.ndindex(na, nb):
    #     for x, y in np.ndindex(nx, ny):
    #         # Extract Choi operators for given x,a and y,b
    #         # Assuming ChoiA_xa[x][a] and ChoiB_yb[y][b] are jnp arrays
    #         AB = jax_kron(Choi_xa[x][a], Choi_yb[y][b])
    #         val = jnp.real(jnp.trace(W @ AB))
    #         prob_abxy = prob_abxy.at[a, b, x, y].set(val)
            
    # define the “trace of Kron” scalar function
    f_AB = lambda A, B: jnp.real(jnp.trace(W @ jax_kron(A, B)))

    # 1) map over B (axis-1 of Choi_yb), for a fixed A:
    f_b   = jax.vmap(f_AB, in_axes=(None, 0))   # (A, [B₀,B₁,…,B_{nb-1}]) → [val₀,…]

    # 2) map that over A (axis-1 of Choi_xa):
    f_ab  = jax.vmap(f_b,   in_axes=(0, None))   # ([A₀,…], B) → matrix (na, nb)

    # 3) then map over y (axis-0 of Choi_yb):
    f_aby = jax.vmap(f_ab,  in_axes=(None, 0))   # (A, [yb₀,…,yb_{ny-1}]) → (ny, na, nb)

    # 4) finally map over x (axis-0 of Choi_xa):
    f_abxy = jax.vmap(f_aby, in_axes=(0, None))  # ([xa₀,…], YB) → (nx, ny, na, nb)

    # run it:
    prob = f_abxy(Choi_xa, Choi_yb)              # shape (nx, ny, na, nb)

    # reorder to your (na, nb, nx, ny):
    prob_abxy = jnp.transpose(prob, (2, 3, 0, 1))       

    penalty = penalty_A + penalty_B + penalty_W

    return prob_abxy, penalty



@partial(jax.jit, static_argnums=(1, 2, 3, 4, 5, 6, 7, 8))
def give_probability_rank1_choi(theta, na, nb, nx, ny, dim_A_in, dim_A_out, dim_B_in, dim_B_out, 
                                unitary_generators_A, unitary_generators_B, 
                                ini_rank1_Choi_A, ini_rank1_Choi_B, 
                                nof_unitary_generators_A, nof_unitary_generators_B):
    nof_params_A = nof_unitary_generators_A * nx
    nof_params_B = nof_unitary_generators_B * ny
    nof_params_W = (dim_A_in * dim_A_out * dim_B_in * dim_B_out)**2
    
    params_A = theta[:nof_params_A]
    params_B = theta[nof_params_A:nof_params_A+nof_params_B]
    params_W = theta[nof_params_A+nof_params_B:]
    
    Choi_x, penalty_A = give_rank1_Choi(params_A, dim_A_in, dim_A_out, na, nx, ini_rank1_Choi_A, unitary_generators_A, nof_unitary_generators_A)
    Choi_y, penalty_B = give_rank1_Choi(params_B, dim_B_in, dim_B_out, nb, ny, ini_rank1_Choi_B, unitary_generators_B, nof_unitary_generators_B)
    W, penalty_W = give_processmatrix(params_W, (dim_A_in, dim_A_out), (dim_B_in, dim_B_out))
    
    Choi_xa = extract_Choi_xa_all(Choi_x, dim_A_in, dim_A_out, na)
    Choi_yb = extract_Choi_xa_all(Choi_y, dim_B_in, dim_B_out, nb)
    
    # # Vectorize over outcomes and settings for party A
    # def compute_A_block(xa):
    #     # xa has shape (na, dimA, dimA)
    #     # vmap over a to produce a function that given a particular Choi_yb and W produces a full slice
    #     return jax.vmap(lambda A: jax.vmap(
    #         lambda YB: jax.vmap(lambda yb: jnp.real(jnp.trace(W @ jax_kron(A, yb))))(YB),
    #         in_axes=(None,)
    #     ))(Choi_yb)
    
    # prob_abxy = jax.vmap(compute_A_block, in_axes=(0,))(Choi_xa)
    # # prob_abxy now has shape (nx, na, nb, ny)
    # # reshape to (na, nb, nx, ny)
    # prob_abxy = jnp.transpose(prob_abxy, (1, 2, 0, 3))
    
    # Create a probability array
    # shape is (na, nb, nx, ny)
    prob_abxy = jnp.zeros((na, nb, nx, ny), dtype=jnp.float32)
    for a, b in np.ndindex(na, nb):
        for x, y in np.ndindex(nx, ny):
            # Extract Choi operators for given x,a and y,b
            # Assuming ChoiA_xa[x][a] and ChoiB_yb[y][b] are jnp arrays
            AB = jax_kron(Choi_xa[x][a], Choi_yb[y][b])
            val = jnp.real(jnp.trace(W @ AB))
            prob_abxy = prob_abxy.at[a, b, x, y].set(val)

    penalty = penalty_A + penalty_B + penalty_W

    return prob_abxy, penalty
    
@partial(jax.jit, static_argnums=(2, 3, 4, 5, 6, 7, 8, 9))
def causalineq_value_rank1(theta, ineq_array, na, nb, nx, ny, dim_A_in, dim_A_out, dim_B_in, dim_B_out,
                           unitary_generators_A, unitary_generators_B, 
                                ini_rank1_Choi_A, ini_rank1_Choi_B, 
                                nof_unitary_generators_A, nof_unitary_generators_B):
    prob_abxy, penalty = give_probability_rank1_choi(theta, na, nb, nx, ny, dim_A_in, dim_A_out, dim_B_in, dim_B_out, 
                                unitary_generators_A, unitary_generators_B, 
                                ini_rank1_Choi_A, ini_rank1_Choi_B, 
                                nof_unitary_generators_A, nof_unitary_generators_B)
    return jnp.sum(ineq_array * prob_abxy), penalty

@partial(jax.jit, static_argnums=(2, 3, 4, 5, 6, 7, 8, 9))
def causalineq_value(theta, ineq_array, na, nb, nx, ny, dim_A_in, dim_A_out, dim_B_in, dim_B_out):
    prob_abxy, penalty = give_probability(theta, na, nb, nx, ny, dim_A_in, dim_A_out, dim_B_in, dim_B_out)
    return jnp.sum(ineq_array * prob_abxy), penalty



# def build_params_for_desired_output(
#     rng_key,
#     desired_out,
#     M,
# ):
#     """
#     Build an M-layer (N->N) network's parameters so that 
#     forward_mlp(params, x0) = desired_out (in [0,1]^N).
    
#     Each layer is NxN. M-1 layers = ReLU, final layer = ReLU + (1 - e^-).
#     The approach:
#       1) Let x_M = desired_out.
#       2) Invert x_{M-1} = -ln(1 - x_M).
#       3) For i in M-2..0:
#          - randomly choose x_i >=0
#          - randomly choose W_i
#          - define b_i = x_{i+1} - W_i x_i
#       4) For i=M-1:
#          - we already have x_{M-1}, x_{M}
#          - randomly choose W_{M-1}
#          - define b_{M-1} = x_{M-1} - W_{M-1} x_{M-2}
    
#     Returns:
#       params: list of length M, each is (W_i, b_i).
#       x_layers: list of x_i for i=0..M, so you can inspect them.
#     """
#     N = desired_out.shape[0]  # dimension
#     # We'll store x_i for i=0..M
#     x_layers = [None]*(M+2)

#     # 1) x_0 is given
#     x_layers[0] = jnp.array(1.0).reshape((1,)) # shape (N,)

#     # 2) x_M = desired_out
#     x_layers[M+1] = desired_out  # shape (N,)
#     # invert final: x_{M-1} = -ln(1 - x_M)
#     # but we do this after we define x_{M-1}, see below

#     # We'll define x_{M-1} from the last layer: 
#     # x_{M} = 1 - e^- ( W_{M-1} x_{M-2} + b_{M-1} ).
#     # We want x_{M-1} = W_{M-1} x_{M-2} + b_{M-1} >=0, 
#     # and x_M = 1 - e^{-x_{M-1}}. So x_{M-1} = -ln(1 - x_M).

#     # But we must define x_{M-2} first. Then define random W_{M-1}, etc.

#     # So let's define x_1.. x_{M-2} as random nonnegative, then define x_{M-1} from desired_out.

#     current_key = rng_key
#     # For layers i=1..M-2, define x_i
#     for i in reversed(range(1, M)):
#         current_key, subkey = random.split(current_key)
#         # random nonnegative (N,)
#         # e.g. take absolute value of normal
#         x_i = jnp.abs(random.normal(subkey, (N,)))
#         x_layers[i] = x_i

#     # define x_{M-1} from desired_out
#     x_layers[M] = -jnp.log(jnp.clip(1.0 - x_layers[M+1], 1e-12, None)) 
#     # ensure it doesn't blow up if x_M=1

#     # Now build W_i, b_i for i=0..M-2, ensuring x_{i+1} = W_i x_i + b_i
#     params = []
#     for i in reversed(range(M)):
#         x_i = x_layers[i]
#         x_ip1 = x_layers[i+1]
#         current_key, subkey = random.split(current_key)
#         # random NxN
#         if i != 0:
#             W_i = random.normal(subkey, (N, N)) * 0.01
#         else:
#             W_i = random.normal(subkey, (N, 1)) * 0.01
        
#         # Solve b_i = x_{i+1} - W_i x_i
#         b_i = x_ip1 - (W_i @ x_i)
#         params.append((W_i, b_i))

#     # Last layer i=M-1
#     # x_{M-1} -> x_M = 1 - e^{-ReLU(W_{M-1} x_{M-2}+ b_{M-1})}
#     # But we've forced ReLU to be trivial by x_{M-1} >=0
#     # so x_{M-1} = W_{M-1} x_{M-2} + b_{M-1}
#     i = M
#     x_im1 = x_layers[i-1]
#     x_i = x_layers[i]  # x_{M-1}
#     current_key, subkey = random.split(current_key)
#     W_i = random.normal(subkey, (N, N)) * 0.01
#     b_i = x_i - (W_i @ x_im1)
#     params.append((W_i, b_i))
#     params.append((x_i.reshape((1,x_i.shape[0])), 0))
    
    

#     return list(reversed(params)), x_layers

# @jax.jit
# def forward_mlp(params, x):
#     """
#     Forward pass of the MLP with M layers, each NxN.
    
#     Args:
#       params: list of (W, b) pairs, length = M
#       x: jnp array of shape (..., N)
#          (Can be a single sample or a batch, as long as last dim = N)
#       N: static integer dimension of each layer
    
#     Returns:
#       jnp array of shape (..., N), the final output.
      
#     For demonstration, we'll do:
#       - ReLU for the first M-1 layers
#       - Sigmoid for the last layer
#     """
#     M = len(params)
    
#     x = jnp.array(x).reshape((1,))
    
#     # For the first M-1 layers: Dense NxN + ReLU
#     for i in range(M - 1):
#         W, b = params[i]
#         x = jnp.dot(x, W) + b
#         x = jnp.maximum(x, 0)  # ReLU
    
#     # Last layer: Dense NxN + Sigmoid
#     W_final, b_final = params[-1]
#     x = jnp.dot(x, W_final) + b_final
#     x = jnp.maximum(x, 0)  # ReLU
#     x = 1 - jnp.exp(-x) # Maps [0,+inf) to [0,1]
#     return x

@jax.jit
def map_0inf_to_01(x):
    # Option A: 1 - exp(-x)
    return 1.0 - jnp.exp(-x)

@jax.jit
def map_0inf_to_01_alt(x):
    # Option B: x / (1 + x)
    return x / (1.0 + x)

@jax.jit
def forward_mlp(params, x0):
    """
    Forward pass of an M-layer network, each layer NxN (except layer 1 is Nx1).
    Each layer uses ReLU, final activation = 1 - e^{-x}.

    params is a list of (W, b).
     - params[0]: W shape (N,1), b shape (N,)
     - params[i>0]: W shape (N,N), b shape (N,)
    x0: shape (1,), the 'trivial input' [1.0]
    Returns: shape (N,) in [0,1].
    """
    x = x0
    for (W, b) in params:
        x = W @ x + b
        x = jnp.maximum(x, 0)  # ReLU
    # final activation
    x = map_0inf_to_01(x)
    # x = map_0inf_to_01_alt(x)
    return x

def build_params_for_desired_output(rng, desired_out, M):
    """
    Build an M-layer feedforward network that yields 'desired_out' in [0,1]^N 
    at input x0=[1.0]. The final activation is 1 - e^{-x}.

    Steps:
      1) x_M = -log(1 - desired_out) >= 0
      2) For i in [M..2], pick random x_{i-1} >= 0, random W_i => b_i = x_i - W_i x_{i-1}
      3) For i=1, W_1 shape=(N,1), x_0=[1], b_1 = x_1 - W_1 x_0

    Returns:
      params: list of length M => (W, b)
      x_layers: [x_0,..., x_M]
    """
    N = desired_out.size
    # Invert final layer
    x_M = -jnp.log(jnp.clip(1.0 - desired_out, 1e-12, 1.0))

    x_layers = [None]*(M+1)
    x_layers[M] = x_M

    # random hidden vectors x_{i-1} for i=M..2
    key = rng
    for i in range(M, 1, -1):
        key, subkey = random.split(key)
        x_im1 = jnp.abs(random.normal(subkey, (N,)))
        x_layers[i-1] = x_im1

    # x_0 = [1.0]
    x_layers[0] = jnp.array([1.0], dtype=jnp.float32)

    # Build (W,b) for each layer
    params = []
    current_key = key

    # Layer 1: (W_1 shape=(N,1), b_1 shape=(N,))
    x_0 = x_layers[0]
    x_1 = x_layers[1]
    current_key, subkey = random.split(current_key)
    W_1 = 0.01 * random.normal(subkey, (N,1))
    b_1 = x_1 - jnp.dot(W_1, x_0)
    params.append((W_1, b_1))

    # Layers 2..M: (N,N)
    for i in range(2, M+1):
        x_im1 = x_layers[i-1]
        x_i = x_layers[i]
        current_key, subkey = random.split(current_key)
        W_i = 0.01 * random.normal(subkey, (N,N))
        b_i = x_i - jnp.dot(W_i, x_im1)
        params.append((W_i, b_i))

    return params, x_layers


# @partial(jax.jit, static_argnums=(2, 3, 4, 5, 6, 7, 8, 9, 10))
# def loss_function(theta, ineq_array, na, nb, nx, ny, dim_A_in, dim_A_out, dim_B_in, dim_B_out, penalty_weight):
#     value, penalty = causalineq_value(theta, ineq_array, na, nb, nx, ny, dim_A_in, dim_A_out, dim_B_in, dim_B_out)
#     return value - penalty_weight*penalty/3

@partial(jax.jit, static_argnums=(3, 4, 5, 6, 7, 8, 9, 10, 11))
def loss_function(theta, ineq_array, mlp_input, na, nb, nx, ny, dim_A_in, dim_A_out, dim_B_in, dim_B_out, penalty_weight):
    theta_physical = forward_mlp(theta, mlp_input)
    value, penalty = causalineq_value(theta_physical, ineq_array, na, nb, nx, ny, dim_A_in, dim_A_out, dim_B_in, dim_B_out)
    return value - penalty_weight*penalty/3

@partial(jax.jit, static_argnums=(3, 4, 5, 6, 7, 8, 9, 10, 11))
def loss_function_rank1(theta, ineq_array, mlp_input, na, nb, nx, ny, dim_A_in, dim_A_out, dim_B_in, dim_B_out, penalty_weight, unitary_generators_A, unitary_generators_B, 
                                ini_rank1_Choi_A, ini_rank1_Choi_B, 
                                nof_unitary_generators_A, nof_unitary_generators_B):
    theta_physical = forward_mlp(theta, mlp_input)
    value, penalty = causalineq_value(theta_physical, ineq_array, na, nb, nx, ny, dim_A_in, dim_A_out, dim_B_in, dim_B_out,
                           unitary_generators_A, unitary_generators_B, 
                                ini_rank1_Choi_A, ini_rank1_Choi_B, 
                                nof_unitary_generators_A, nof_unitary_generators_B)
    return value - penalty_weight*penalty/3


# %%
def is_choi_cptp(choi, dimX, dimY, choi_convention='R', verbose=True, tol=1e-5):
    """
    Check if the Choi matrix of `C: X -> Y` is valid. 
    
    If `choi_convention='R'`:
    choi = sum_{i,j} |i><j|_X \otimes C(|i><j|_X)
    
    If `choi_convention='L'2`:
    choi = sum_{i,j} C(|i><j|_X) \otimes |i><j|_X.
    """
    if dimX == 1:
        minimum_eigenvalue = np.math.real(np.linalg.eigvalsh(np.reshape(choi, [dimY, dimY]))[0])
        if minimum_eigenvalue < -tol:
            if verbose:
                print('Minimum eigenvalue of Choi matrix is', minimum_eigenvalue)
            return False
        # Trace 1
        trace = np.math.real(np.linalg.trace(np.reshape(choi, [dimY, dimY])))
        if not np.isclose(trace, 1, atol=tol):
            if verbose:
                print(f'Trace of Choi matrix is not 1 but {trace}')
            return False
        return True    
    
    _dtype = choi.dtype
    
    _reshaped = np.reshape(choi, [dimX*dimY, dimX*dimY])
    choi_H = (_reshaped + _reshaped.conj().T)/2
    diff = np.sum(np.abs(np.reshape(choi_H - _reshaped, [-1]))**2)
    if diff > tol:
        if verbose:
            print(f'Choi matrix is not Hermitian. diff={diff}')
        return False
    
    minimum_eigenvalue = np.real(np.linalg.eigvalsh(_reshaped)[0])
    if minimum_eigenvalue < -tol:
        if verbose:
            print('Minimum eigenvalue of Choi matrix is', minimum_eigenvalue)
        return False
    
    if choi_convention == 'R':
        trY = partial_trace(choi, (dimX, dimY), (1,))
    elif choi_convention == 'L':
        trY = partial_trace(choi, (dimY, dimX), (0,))
    else:
        raise ValueError('Invalid choi_convention')
    
    diff = np.sum(np.abs(np.reshape(trY - np.eye(dimX, dtype=_dtype),#/dimX,
                                                [-1]))**2)
    if diff > tol:
        if verbose:
            print(f'Partial trace of Choi matrix is not identity. diff_norm={diff} {trY}')
        return False
    return True


def is_choi_2party_processmatrix(W, dim_A_in, dim_A_out, dim_B_in, dim_B_out, tol=1e-5, verbose=True):
    def _Lv(W):
        # 2T + 4T − 24T − 34T + 234T − 12T + 124T
        W_reshaped = np.reshape(W, (dim_A_in, dim_A_out, dim_B_in, dim_B_out,)*2)
        AO_W = np_partial_trace_with_id(W_reshaped, [1])
        BO_W = np_partial_trace_with_id(W_reshaped, [3])
        AOBO_W = np_partial_trace_with_id(W_reshaped, [1, 3])
        BIBO_W = np_partial_trace_with_id(W_reshaped, [2, 3])
        AOBIBO_W = np_partial_trace_with_id(W_reshaped, [1, 2, 3])
        AIAO_W = np_partial_trace_with_id(W_reshaped, [0, 1])
        AIAOBO_W = np_partial_trace_with_id(W_reshaped, [0, 1, 3])
        W_projected = AO_W + BO_W - AOBO_W - BIBO_W  + AOBIBO_W - AIAO_W + AIAOBO_W
        return W_projected
    
    projected = _Lv(W)
    
    if not np.allclose(projected, W, atol=tol):
        if verbose:
            print(f"The process matrix is not in the correct linear subspace, but is {np.linalg.norm(projected-W)} away.")
        return False
    
    mineigW = np.linalg.eigvalsh(W)
    if mineigW[0] < -tol:
        if verbose:
            print(f"The process matrix has negative eigenvalues. {mineigW[0]}")
        return False
    
    if not np.isclose(np.trace(W), dim_A_out*dim_B_out, atol=tol):
        if verbose:
            print(f"The process matrix does not have trace dim_A_out*dim_B_out={dim_A_out*dim_B_out} but {np.trace(W)}.")
        return False
        
    return True

def is_choi_3party_processmatrix(W, dim_A_in, dim_A_out, dim_B_in, dim_B_out, dim_C_in, dim_C_out, tol=1e-5, verbose=True):
    def _Lv_3party(W):
        _dims = (dim_A_in, dim_A_out, dim_B_in, dim_B_out, dim_C_in, dim_C_out)
        W = np.reshape(W, _dims*2)
    
        AIAOBIBOCI_W = np_partial_trace_with_id(W, (0, 1, 2, 3, 4))
        AIAOBIBO_W   = np_partial_trace_with_id(W, (0, 1, 2, 3))
        AIAOBOCICO_W = np_partial_trace_with_id(W, (0, 1, 3, 4, 5))
        AIAOBOCI_W   = np_partial_trace_with_id(W, (0, 1, 3, 4))
        AIAOBO_W     = np_partial_trace_with_id(W, (0, 1, 3))
        
        # second line in the comment above
        AIAOCICO_W   = np_partial_trace_with_id(W, (0, 1, 4, 5))
        AIAOCI_W     = np_partial_trace_with_id(W, (0, 1, 4))
        AIAO_W       = np_partial_trace_with_id(W, (0, 1))
        AOBIBOCICO_W = np_partial_trace_with_id(W, (1, 2, 3, 4, 5))
        AOBIBOCI_W   = np_partial_trace_with_id(W, (1, 2, 3, 4))
        AOBIBO_W     = np_partial_trace_with_id(W, (1, 2, 3))
        
        # third line in the comment above
        AOBOCICO_W   = np_partial_trace_with_id(W, (1, 3, 4, 5))
        AOBOCI_W     = np_partial_trace_with_id(W, (1, 3, 4))
        AOBO_W       = np_partial_trace_with_id(W, (1, 3))
        AOCICO_W     = np_partial_trace_with_id(W, (1, 4, 5))
        AOCI_W       = np_partial_trace_with_id(W, (1, 4))
        AO_W         = np_partial_trace_with_id(W, (1,))
        BIBOCICO_W   = np_partial_trace_with_id(W, (2, 3, 4, 5))
        
        # fourth line in the comment above
        BIBOCI_W     = np_partial_trace_with_id(W, (2, 3, 4))
        BIBO_W       = np_partial_trace_with_id(W, (2, 3))
        BOCICO_W     = np_partial_trace_with_id(W, (3, 4, 5))
        BOCI_W       = np_partial_trace_with_id(W, (3, 4))
        BO_W         = np_partial_trace_with_id(W, (3,))
        CICO_W       = np_partial_trace_with_id(W, (4, 5))
        CI_W         = np_partial_trace_with_id(W, (4,))

        # Apply the projection
        W_projected = (AIAOBIBOCI_W - AIAOBIBO_W + AIAOBOCICO_W - AIAOBOCI_W + 
                    AIAOBO_W - AIAOCICO_W + AIAOCI_W - AIAO_W + AOBIBOCICO_W - 
                    AOBIBOCI_W + AOBIBO_W - AOBOCICO_W + AOBOCI_W - AOBO_W + 
                    AOCICO_W - AOCI_W + AO_W - BIBOCICO_W + BIBOCI_W - BIBO_W + 
                    BOCICO_W - BOCI_W + BO_W - CICO_W + CI_W)
        
        W_projected = (W_projected + W_projected.conj().T)/2
        
        return W_projected
    
    projected = _Lv_3party(W)
    
    if not np.allclose(projected, W, atol=tol):
        if verbose:
            print(f"The process matrix is not in the correct linear subspace, but is {np.linalg.norm(projected-W)} away.")
        return False
    
    mineigW = np.linalg.eigvalsh(W)
    if mineigW[0] < -tol:
        if verbose:
            print(f"The process matrix has negative eigenvalues. {mineigW[0]}")
        return False
    
    if not np.isclose(np.trace(W), dim_A_out*dim_B_out, atol=tol):
        if verbose:
            print(f"The process matrix does not have trace dim_A_out*dim_B_out={dim_A_out*dim_B_out} but {np.trace(W)}.")
        return False
        
    return True

def np_partial_trace(tensor, keep_indices):
    """Takes the partial trace of a given tensor.

    The input tensor must have shape `(d_0, ..., d_{k-1}, d_0, ..., d_{k-1})`.
    The trace is done over all indices that are not in keep_indices. The
    resulting tensor has shape `(d_{i_0}, ..., d_{i_r}, d_{i_0}, ..., d_{i_r})`
    where `i_j` is the `j`th element of `keep_indices`.

    Args:
        tensor: The tensor to sum over. This tensor must have a shape
            `(d_0, ..., d_{k-1}, d_0, ..., d_{k-1})`.
        keep_indices: Which indices to not sum over. These are only the indices
            of the first half of the tensors indices (i.e. all elements must
            be between `0` and `tensor.ndims / 2 - 1` inclusive).

    Raises:
        ValueError: if the tensor is not of the correct shape or the indices
            are not from the first half of valid indices for the tensor.
            
    # SOURCE: Cirq code https://github.com/quantumlib/Cirq/blob/v1.4.0/cirq-core/cirq/linalg/transformations.py#L382-L417
    # but modified to use tf.einsum with string equation
    """
    ndim = tensor.ndim // 2
    if not all(tensor.shape[i] == tensor.shape[i + ndim] for i in range(ndim)):
        raise ValueError(
            f'Tensors must have shape (d_0,...,d_{{k-1}},d_0,...,'
            f'd_{{k-1}}) but had shape ({tensor.shape}).'
        )
    if not all(i < ndim for i in keep_indices):
        raise ValueError(
            f'keep_indices were {keep_indices} but must be in first half, '
            f'i.e. have index less that {ndim}.'
        )
    keep_set = set(keep_indices)
    keep_map = dict(zip(keep_indices, sorted(keep_indices)))
    left_indices = [keep_map[i] if i in keep_set else i for i in range(ndim)]
    right_indices = [ndim + i if i in keep_set else i for i in left_indices]
    final_indices = left_indices + right_indices
    index_to_letter = dict(zip(range(ndim * 2), ''.join([chr(ord('a') + i) for i in range(ndim*2)])))
    str_eq = ''.join([index_to_letter[i] for i in final_indices])
    str_eq += '->'
    # find in final_indices indices that appear only once
    # and add them to the output string
    for i in range(ndim * 2):
        if final_indices.count(i) == 1:
            str_eq += index_to_letter[i]
    return np.einsum(str_eq, tensor)

def np_swap_subspaces(state, dims, perm):
    _perm = [*perm, *(len(perm) + np.array(perm)).tolist()]
    _dims = [*dims, *dims]
    state = np.transpose(np.reshape(state, _dims), axes=_perm)
    return np.reshape(state, (np.prod(dims, dtype=int),)*2)

def np_partial_trace_with_id(tensor, remove_indices):
    from utils import find_permutation
    ndim = tensor.ndim // 2
    all_indices = list(range(ndim))
    keep_indices = [i for i in all_indices if i not in remove_indices]

    dims_of_removed = [tensor.shape[i] for i in remove_indices]
    dims_of_kept = [tensor.shape[i] for i in keep_indices]
    
    _tensor = tensor.copy()
    
    _tensor = np_partial_trace(_tensor.copy(), keep_indices)
    normalized_identity = np.eye(np.prod(dims_of_removed, dtype=int)) / np.prod(dims_of_removed, dtype=int)
    _tensor = np.kron(np.reshape(_tensor, (np.prod(dims_of_kept, dtype=int),)*2), normalized_identity)
        
    perm = find_permutation(keep_indices + remove_indices, all_indices)
    
    final = np_swap_subspaces(_tensor, dims_of_kept + dims_of_removed, perm)
    
    return final


def np_make_hermitian(A):
    return (A + A.conj().T) / 2

def np_random_hermitian(n):
    A = np.random.randn(n, n) + 1j * np.random.randn(n, n)
    return np_make_hermitian(A)

def unitary2choi(U, basis, d):
    C = 0
    for i in range(d**2):
        C += np.kron(basis[i, :, :], U @ basis[i, :, :] @ U.conj().T)
    return C

def generate_random_choi(dim_in, dim_out, dtype=np.complex128):
    dim = dim_in * dim_out
    correct_tr_value = dim_in
    C = np_random_hermitian(dim_in*dim_out).astype(dtype)
    C = Lv_projector_channel(C, dim_in, dim_out)
    C = (C + C.conj().T) / 2
    C = correct_tr_value * (C/np.trace(C))
    
    if not np.isclose(np.trace(C), correct_tr_value):
        print(f"generate_random_choi: Projection modified the trace from {correct_tr_value} to {np.trace(C)}")

    mineig = np.min(np.linalg.eigvalsh(C))
    if mineig < 1e-7:
        eps = mineig
        n = dim  / correct_tr_value
        p = (1e-6 - 1/n) / (eps - 1/n)
        C = p * C + (1-p) * 1/n * np.eye(dim)
    
    is_choi_cptp(C, dim_in, dim_out, verbose=True)

    return C

def generate_random_2party_process(dim_A_in, dim_B_in, dim_A_out, dim_B_out, dtype=np.complex128):
    dim = dim_A_in * dim_B_in * dim_A_out * dim_B_out
    dim_in = dim_A_out * dim_B_out
    dim_out = dim_A_in * dim_B_in
    correct_tr_value = dim_A_out * dim_B_out
    C = np_random_hermitian(dim_in*dim_out).astype(dtype)
   
    C = Lv_projector_process_2parties(C, dim_A_in, dim_A_out, dim_B_in, dim_B_out)
    C = (C + C.conj().T) / 2
    C = correct_tr_value * (C/np.trace(C))

    if not np.isclose(np.trace(C), correct_tr_value):
        print(f"generate_random_2party_process: Projection modified the trace from {correct_tr_value} to {np.trace(C)}")

    mineig = np.min(np.linalg.eigvalsh(C))
    if mineig < 1e-7:
        eps = mineig
        n = dim  / correct_tr_value
        p = (1e-6 - 1/n) / (eps - 1/n)
        C = p * C + (1-p) * 1/n * np.eye(dim)
        
    is_choi_2party_processmatrix(C, dim_A_in, dim_A_out, dim_B_in, dim_B_out, verbose=True)

    return C

def generate_random_3party_process(dim_A_in, dim_B_in, dim_A_out, dim_B_out, dim_C_in, dim_C_out, dtype=np.complex128):
    dim = dim_A_in * dim_B_in * dim_A_out * dim_B_out * dim_C_in * dim_C_out
    dim_in = dim_A_out * dim_B_out * dim_C_out
    dim_out = dim_A_in * dim_B_in * dim_C_in
    correct_tr_value = dim_A_out * dim_B_out * dim_C_out
    
    C = np_random_hermitian(dim_in*dim_out).astype(dtype)
   
    C = Lv_projector_process_3parties(C, dim_A_in, dim_A_out, dim_B_in, dim_B_out, dim_C_in, dim_C_out)
    C = (C + C.conj().T) / 2
    C = correct_tr_value * (C/np.trace(C))

    if not np.isclose(np.trace(C), correct_tr_value):
        print(f"generate_random_2party_process: Projection modified the trace from {correct_tr_value} to {np.trace(C)}")

    mineig = np.min(np.linalg.eigvalsh(C))
    if mineig < 1e-7:
        eps = mineig
        n = dim  / correct_tr_value
        p = (1e-6 - 1/n) / (eps - 1/n)
        C = p * C + (1-p) * 1/n * np.eye(dim)
        
    is_choi_3party_processmatrix(C, dim_A_in, dim_A_out, dim_B_in, dim_B_out, dim_C_in, dim_C_out, verbose=True)

    return C

def generate_random_Choi_x(dim_in, dim_out,  nof_outcomes, nof_settings):
    return np.stack([generate_random_choi(dim_in, dim_out*nof_outcomes) for _ in range(nof_settings)], axis=0)

def choi_x2theta(Choi_x):
    theta = [hermitian_to_param01(Choi_x[x]) for x in range(Choi_x.shape[0])]
    return np.concatenate(theta)

def choi_x_process2theta(Choi_x, Choi_y, W):
    theta = [hermitian_to_param01(Choi_x[x]) for x in range(Choi_x.shape[0])]
    theta += [hermitian_to_param01(Choi_y[y]) for y in range(Choi_y.shape[0])]
    theta += [hermitian_to_param01(W)]
    return np.concatenate(theta)

# %%
def test_Choi_x_inverse_theta():
    Choi_x = generate_random_Choi_x(2, 2, 2, 3)
    theta = choi_x2theta(Choi_x)
    Choi_x2, _ = give_Choi(jnp.array(theta), 2, 2, 2, 3)
    if np.allclose(Choi_x2, Choi_x):
        return True
    else:
        return False
    
# print("Test passed:", test_Choi_x_inverse_theta())

def test_process_inverse_theta():
    W = generate_random_2party_process(2, 2, 2, 2)
    theta = hermitian_to_param01(W)
    W2, _ = give_processmatrix(jnp.array(theta), (2, 2), (2, 2))
    if np.allclose(W2, W):
        return True
    else:
        return False
    
# print("Test passed:", test_process_inverse_theta())

def delta_ij(i, j):
    return 1 if i == j else 0

def generate_ops(outcomes_per_party, settings_per_party):
    """generate symbolic operators of the form partyString_input_output
    that can be used to construct bell inequalities over which we can
    optimise"""
    import sympy as sp
    ops = []
    for p, nx in settings_per_party.items():
        ops_x = []
        for x in range(nx):
            ops_x.append([sp.Symbol(f'{p}_{x}_{a}')
                          for a in range(outcomes_per_party[p])])
        ops.append(ops_x)
    return ops

def Bell_CG2prob(belloperator, outcomes_per_party: dict, settings_per_party: dict):
    """
    Takes a Bell operator in CG and returns it in probability space.
    """

    import numpy as np
    import sympy as sp
    
    nr_parties = len(outcomes_per_party)
    parties = sorted(list(outcomes_per_party.keys()))
    bell_prob = np.zeros((*[outcomes_per_party[p] for p in parties],
                          *[settings_per_party[p] for p in parties]))
    sym_expanded_constant_term = 0
    template_const = np.zeros((*[outcomes_per_party[p] for p in parties],
                               *[settings_per_party[p] for p in parties]))  
    for outcomes in np.ndindex(*[outcomes_per_party[p] for p in parties]):
        sym_expanded_constant_term += np.prod([sp.symbols(f"{p}_0_{outcomes[i]}")
                                               for i, p in enumerate(parties)])
        template_const[(*outcomes,*(0,)*nr_parties)] += 1
        
    expanded_bell = 0
    belloperator = sp.expand(belloperator)
    for term, coeff in belloperator.as_coefficients_dict().items():
        # if term == 1:
        #     bell_prob += float(coeff)*template_const
        #     expanded_bell += coeff*sym_expanded_constant_term
        #     # expanded_bell += coeff
        # else:
        _, ops = term.as_coeff_mul()
        if len(ops) == nr_parties:
            outs = [int(n.split('_')[-1]) for n in str(term).split('*')]
            ins  = [int(n.split('_')[-2]) for n in str(term).split('*')]
            bell_prob[(*outs, *ins)] += float(coeff)
            expanded_bell += coeff*term
        else:
            # Add the missing operators
            present_parties = [n.split('_')[0] for n in str(term).split('*')]
            missing_parties = [p for p in parties if p not in present_parties] if present_parties else []
            expanded_term = term
            for p in missing_parties:
                expanded_term *= sum([sp.Symbol(f'{p}_0_{a}')
                                    for a in range(outcomes_per_party[p])])
            expanded_term = sp.expand(expanded_term)
            expanded_bell += coeff*expanded_term
            for term2, coeff2 in expanded_term.as_coefficients_dict().items():
                outs = [int(n.split('_')[-1]) for n in str(term2).split('*')]
                ins  = [int(n.split('_')[-2]) for n in str(term2).split('*')]
                bell_prob[(*outs, *ins)] += float(coeff)*float(coeff2)

    return bell_prob


if __name__ == "__main__":
    outcomes_per_party = {'A': 2, 'B': 2}
    settings_per_party = {'A': 2, 'B': 2}
    ops = generate_ops(outcomes_per_party, settings_per_party)
    p_ab_xy = np.zeros((outcomes_per_party['A'], outcomes_per_party['B'],
                        settings_per_party['A'], settings_per_party['B']), dtype=object)
    for x, y in np.ndindex(settings_per_party['A'], settings_per_party['B']):
        for a, b in np.ndindex(outcomes_per_party['A'], outcomes_per_party['B']):
            p_ab_xy[a, b, x, y] = ops[0][x][a] * ops[1][y][b]
            
    GYNI = 0
    for x, y in np.ndindex(settings_per_party['A'], settings_per_party['B']):
        for a, b in np.ndindex(outcomes_per_party['A'], outcomes_per_party['B']):
            GYNI += delta_ij(a, y) * delta_ij(b, x) * p_ab_xy[a, b, x, y]
    GYNI = GYNI / 4

    GYNI_as_array = Bell_CG2prob(GYNI, outcomes_per_party, settings_per_party)
    ineq_array = np.array(GYNI_as_array).astype(np.float32)

    # LOCAL_DIM = 2
    # dim_A_in = LOCAL_DIM
    # dim_A_out = LOCAL_DIM
    # dim_B_in = LOCAL_DIM
    # dim_B_out = LOCAL_DIM
    dim_A_in = 2
    dim_A_out = 2
    dim_B_in = 2
    dim_B_out = 2
    na = outcomes_per_party['A']
    nb = outcomes_per_party['B']
    nx = settings_per_party['A']
    ny = settings_per_party['B']
    
    # filename = f"/Users/emi/Documents/Code/indefcausalityML/Choi_subspace_basis_{dim_A_in}_{dim_A_out*na}.npz"
    # subspace_basis, effective_subspace_basis, unitary_generator_basis, orthogonal_basis = \
    #     np.load(filename + '.npz').values()

    # filename_choi = f"./subspace_basis_Choi_dI={dim_A_in}_dO={dim_A_out*na}"
    # subspace_basis, orthogonal_basis, hermsubspace_basis, subspace_basis_filename = \
    #     np.load(filename_choi + '.npz').values()
    # filename_choi_generator = f"./U_generator_basis_Choi_dI={dim_A_in}_dO={dim_A_out*na}"
    # commutator_basis, hermitian_commutator_basis, filename_choi_generator_filename = \
    #     np.load(filename_choi + '.npz').values()


    nof_params_A = (dim_A_in * dim_A_out * na)**2 * nx
    nof_params_B = (dim_B_in * dim_B_out * nb)**2 * ny
    nof_params_W = (dim_A_in * dim_A_out * dim_B_in * dim_B_out)**2
    nof_params_A_rank1 = (dim_A_in * dim_A_out * na)**2 * nx
    nof_params_B_rank1 = (dim_B_in * dim_B_out * nb)**2 * ny
    nof_params_W = (dim_A_in * dim_A_out * dim_B_in * dim_B_out)**2
    # theta = jnp.arange(nof_params_A + nof_params_B + nof_params_W, dtype=jnp.float32)

    Choi_x = generate_random_Choi_x(dim_A_in, dim_A_out, na, nx)
    Choi_y = generate_random_Choi_x(dim_B_in, dim_B_out, nb, ny)
    W = generate_random_2party_process(dim_A_in, dim_B_in, dim_A_out, dim_B_out)
    Choi_xa = extract_Choi_xa_all(Choi_x, dim_A_in, dim_A_out, na)
    Choi_yb = extract_Choi_xa_all(Choi_y, dim_B_in, dim_B_out, nb)
    
    
    print(f"Dim of Choi_xa, Choi_yb, W for local dims {dim_A_in},{dim_A_out},{dim_B_in},{dim_B_out}:",
          Choi_xa[0,0].shape[0], Choi_yb[0,0].shape[0], W.shape[0])
    
    Choi_xa = extract_Choi_xa_all(Choi_x, dim_A_in, dim_A_out, na)
    Choi_yb = extract_Choi_xa_all(Choi_y, dim_B_in, dim_B_out, nb)
    
    
    theta = choi_x_process2theta(Choi_x, Choi_y, W)

    prob, penalty = give_probability(theta, na, nb, nx, ny, dim_A_in, dim_A_out, dim_B_in, dim_B_out)

    _Ca, _pa = give_Choi(theta[:nof_params_A], dim_A_in, dim_A_out, na, nx)
    _Cb, _pb = give_Choi(theta[nof_params_A:nof_params_A+nof_params_B], dim_B_in, dim_B_out, nb, ny)
    _W, _pw = give_processmatrix(theta[nof_params_A+nof_params_B:], (dim_A_in, dim_A_out), (dim_B_in, dim_B_out))

    # print("loss function value:", loss_function(theta, ineq_array, na, nb, nx, ny, dim_A_in, dim_A_out, dim_B_in, dim_B_out, 1e-4))
    # print("eigenvalues of _Ca[0]:", jnp.linalg.eigvalsh(_Ca[0]))
    # print("logdeg of _Ca[0]:", jnp.linalg.slogdet(_Ca[0]))

    # Assume loss_function is defined as above and jitted:
    # loss_function(theta, ineq_array, na, nb, nx, ny, dim_A_in, dim_A_out, dim_B_in, dim_B_out)

    def initialize_random_theta(key, dim_A_in, dim_A_out, na, nb, nx, ny, dim_B_in, dim_B_out):
        # Compute N = number of parameters
        # N = ((dim_A_in * dim_A_out * na)**2 * nx) + ((dim_B_in * dim_B_out * nb)**2 * ny) + ((dim_A_in * dim_A_out * dim_B_in * dim_B_out)**2)
        N = ((dim_A_in * dim_A_out * na)**2 * nx) + ((dim_B_in * dim_B_out * nb)**2 * ny) + ((dim_A_in * dim_A_out * dim_B_in * dim_B_out)**2)
        # Initialize theta randomly, for example with normal distribution
        return jax.random.normal(key, (N,))
    
    @partial(jax.jit, static_argnums=(3,4,5))
    def run_epoch(theta, opt_state, mlp_input, optimizer, steps_per_epoch, penalty_weight):
        """
        Runs a specified number of optimization steps in a JIT-compiled loop.
        
        Args:
            theta: Current parameter vector.
            opt_state: Current optimizer state.
            optimizer: An Optax optimizer object (static).
            steps_per_epoch: Number of steps to run in this epoch (static).
            penalty_weight: A scalar controlling the penalty term (static).
        
        Returns:
            Updated theta and optimizer state.
        """

        def loss_for_theta(p):
            return loss_function(
                p, 
                -ineq_array,  # example
                mlp_input,
                na, nb, nx, ny, 
                dim_A_in, dim_A_out, dim_B_in, dim_B_out, 
                penalty_weight
            )

        grad_fn = jax.grad(loss_for_theta)

        def step_fn(i, carry):
            t, st = carry
            g = grad_fn(t)
            updates, st = optimizer.update(g, st, t)
            t = optax.apply_updates(t, updates)
            return (t, st)
        
        theta, opt_state = lax.fori_loop(0, steps_per_epoch, step_fn, (theta, opt_state))
        return theta, opt_state

    def run_one_training_loop(optimizer, dim_A_in, dim_A_out, dim_B_in, dim_B_out, na, nb, nx, ny, ineq_array, verbose=True):
        num_epochs = 50
        steps_per_epoch = 200
        PENALTY_WEIGHT = 1e-3
        CONVERGENCE_THRESHOLD = 1e-4
        TRIVIAL_INPUT_ARR = jnp.array(1.0).reshape((1,))

        # theta_init = initialize_random_theta(key, dim_A_in, dim_A_out, na, nb, nx, ny, dim_B_in, dim_B_out)
        Choi_x = generate_random_Choi_x(dim_A_in, dim_A_out, na, nx)
        Choi_y = generate_random_Choi_x(dim_B_in, dim_B_out, nb, ny)
        W = generate_random_2party_process(dim_A_in, dim_B_in, dim_A_out, dim_B_out)
        theta_init = choi_x_process2theta(Choi_x, Choi_y, W)
        
        NOF_DEEP_LAYERS = 1
        mlp_weights, x_layers = build_params_for_desired_output(jax.random.PRNGKey(secrets.randbits(32)), theta_init, NOF_DEEP_LAYERS)
        
        theta = forward_mlp(mlp_weights, TRIVIAL_INPUT_ARR)
        if not np.allclose(theta, theta_init):
            print("Initial theta and theta from MLP are not the same.")
        
        last_good_Ca = Choi_x
        last_good_Cb = Choi_y
        last_good_W = W
        last_good_ineq_val = 0

        opt_state = optimizer.init(mlp_weights)
        
        loss_value = loss_function(mlp_weights, -ineq_array, x_layers[0], na, nb, nx, ny, dim_A_in, dim_A_out, dim_B_in, dim_B_out, PENALTY_WEIGHT)
        ineq_val, penalty = causalineq_value(theta_init, ineq_array, na, nb, nx, ny, dim_A_in, dim_A_out, dim_B_in, dim_B_out)

        # Training loop over epochs (Python loop)
        has_converged = False
        avg_epoch_time = 0.0
        for epoch in (pbar := tqdm(range(num_epochs))):
            # Run one epoch (JIT-compiled steps)
            start_time = time.perf_counter()
            mlp_weights, opt_state = jax.block_until_ready(run_epoch(mlp_weights, opt_state, x_layers[0], optimizer, steps_per_epoch, PENALTY_WEIGHT))
            theta = forward_mlp(mlp_weights, x_layers[0])
            end_time = time.perf_counter()
            # avg_epoch_time = (avg_epoch_time * epoch + (end_time - start_time)) / (epoch + 1)
            avg_epoch_time = end_time - start_time
            
            # After the JIT-compiled steps, we can evaluate and print the loss
            new_loss_value = loss_function(mlp_weights, -ineq_array, x_layers[0], na, nb, nx, ny, dim_A_in, dim_A_out, dim_B_in, dim_B_out, PENALTY_WEIGHT)
            ineq_val, penalty = causalineq_value(theta, ineq_array, na, nb, nx, ny, dim_A_in, dim_A_out, dim_B_in, dim_B_out)
            # Check if the Choi channels and process matrix resulting in ineq_val are valid
            _Ca, _pa = give_Choi(theta[:nof_params_A], dim_A_in, dim_A_out, na, nx)
            _Cb, _pb = give_Choi(theta[nof_params_A:nof_params_A+nof_params_B], dim_B_in, dim_B_out, nb, ny)
            _W, _pw = give_processmatrix(theta[nof_params_A+nof_params_B:], (dim_A_in, dim_A_out), (dim_B_in, dim_B_out))
            
            delta_loss = new_loss_value - loss_value
            
            anything_invalid = False
            for x in range(nx):
                if not is_choi_cptp(_Ca[x], dim_A_in, dim_A_out*na):
                    print(f"Choi A {x} is not valid.")
                    anything_invalid = True
            for y in range(ny):
                if not is_choi_cptp(_Cb[y], dim_B_in, dim_B_out*nb):
                    print(f"Choi B {y} is not valid.")
                    anything_invalid = True
            if not is_choi_2party_processmatrix(_W, dim_A_in, dim_A_out, dim_B_in, dim_B_out):
                print("Process matrix not valid.")
                anything_invalid = True
                
            if verbose:
                pbar.set_description(f"Steps: {(epoch+1)*steps_per_epoch}/{steps_per_epoch*num_epochs} Δt: {avg_epoch_time:.2E}, "+
                                     f"Loss: {new_loss_value:.4f} ΔL: {delta_loss:.2E} Cost: {penalty:.2E} f: {ineq_val:.4f}  last good f: {last_good_ineq_val:.4f}")
                
            if anything_invalid:
                print("Invalid Choi or process matrix. Exiting. Returning last good values for ineq_val=", last_good_ineq_val)
                has_converged = False
                break
            else:
                has_converged = True
                last_good_Ca = _Ca
                last_good_Cb = _Cb
                last_good_W = _W
                last_good_ineq_val = ineq_val
                
            if jnp.linalg.norm(delta_loss) < CONVERGENCE_THRESHOLD:
                print("Converged.")
                break
            
            loss_value = new_loss_value
            
        return last_good_ineq_val, last_good_Ca, last_good_Cb, last_good_W, has_converged

    
    def multiple_training_runs(dim_A_in, dim_A_out, dim_B_in, dim_B_out, na, nb, nx, ny, ineq_array,
                            MAX_RUNS=1000,
                            max_iter=100,
                            tol=1e-6,
                            verbose=False):
        """
        Runs the seesaw method multiple times (MAX_RUNS), each time with random
        initial guesses for Cax, Cby, W. Tracks the best final objective value
        found, along with the corresponding best Cax, Cby, W. Finally, writes
        them to file 'best_ineqval_Cax_Cby_W.npz'.

        Args:
            problemA, problemB, problemW: instances of the classes for A/B/W optimization.
            MAX_RUNS: How many random restarts to try.
            max_iter, tol: Parameters passed to the seesaw routine.
            verbose: Whether to print intermediate progress.

        Returns:
            (best_Cax, best_Cby, best_W, best_obj)
            Also saves them in the file 'best_ineqval_Cax_Cby_W.npz'.
        """

        best_obj = -1e100
        best_Cax = None
        best_Cby = None
        best_W   = None

        optimizer = optax.adam(0.0001)
        # optimizer = optax.sgd(learning_rate=0.01, momentum=0.2)

        for run_idx in range(MAX_RUNS):
            if verbose:
                print(f"\n=== Training run #{run_idx+1}/{MAX_RUNS} ===")

            # Call your existing seesaw method
            final_obj, final_Cax, final_Cby, final_W, has_converged = run_one_training_loop(optimizer, dim_A_in, dim_A_out, dim_B_in, dim_B_out, na, nb, nx, ny, ineq_array)
            
            # Keep track of the best objective so far
            if final_obj > best_obj:
                best_obj = final_obj
                best_Cax = final_Cax
                best_Cby = final_Cby
                best_W   = final_W
                
                # After we've done all runs, save the best results to file
                np.savez("mltraining_new_best_ineqval_Cax_Cby_W.npz",
                        best_obj=best_obj,
                        best_Cax=best_Cax,
                        best_Cby=best_Cby,
                        best_W=best_W)
                
            if verbose:
                print(f"  Final objective: {final_obj}.  Best so far is {best_obj}")



        # After we've done all runs, save the best results to file
        np.savez("mltraining_best_ineqval_Cax_Cby_W.npz",
                best_obj=best_obj,
                best_Cax=best_Cax,
                best_Cby=best_Cby,
                best_W=best_W)

        if verbose:
            print("\nAll runs complete.")
            print("Best objective found:", best_obj)
            print("Best results saved in 'best_ineqval_Cax_Cby_W.npz'.")

        return best_Cax, best_Cby, best_W, best_obj
        

    # ineq, Ca, Cb, W = run_one_training_loop(dim_A_in, dim_A_out, dim_B_in, dim_B_out, na, nb, nx, ny, ineq_array)
    multiple_training_runs(dim_A_in, dim_A_out, dim_B_in, dim_B_out, na, nb, nx, ny, ineq_array, verbose=True)