In the process of jaxpr to ir conversion, the problem of parameter mismatch occurs

[a-seeker]

I ran into a problem, because jaxpr doesn't support serialization itself, so I wanted to convert jaxpr to ir to use. In most cases, the conversion of ir works fine, but in a few special cases (I couldn't locate the case), the number of arguments in pr was reduced after the conversion, resulting in parameter mismatch issues and errors in the code. The following is my code, please ask if you have any suggestions to help me solve this problem.

from jax import jit
import jax.numpy as jnp
import os
from functools import partial
import jax.core as jcore
import jax
from jax._src.lib import xla_bridge as xb

os.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=16'


@jit
def complex_fun(*args):
    "function"
    pass

#list of sequence as input
args = []

pr = jax.make_jaxpr(complex_computation)(*args)
f = jit(partial(jcore.eval_jaxpr, pr.jaxpr, []))
print(f)
compiled_ir = str(f.lower(*args).compiler_ir())

print(compiled_ir.as_text())
a = jax.device_put(x)
b = jax.device_put(y)

bakend = xb.get_backend("cuda")

for i in args:
  input.append(jax.device_put(i, device=bakend.local_devices()[device]))
complied = bakend.compile(complied_ir)
output = complied.execute(input)

[hawkinsp]

It's going to be difficult to help you without a runnable reproduction of the problem.