# Buggy code
def xla_callable(fun, *abstract_args):
    pvals = [pe.PartialVal((aval, core.unit)) for aval in abstract_args]
    with core.new_master(pe.JaxprTrace, True) as master:
        jaxpr, (pval, consts, env) = pe.trace_to_subjaxpr(fun, master).call_wrapped(
            pvals
        )
        assert not env  # no subtraces here (though cond might eventually need them)
        compiled, result_shape = compile_jaxpr(jaxpr, consts, *abstract_args)
        del master, consts, jaxpr, env
    handle_result = result_handler(result_shape)
    return partial(execute_compiled, compiled, pval, handle_result)


# Traceback (RuntimeError)
$ python loop.py
/usr/local/google/home/phawkins/p/jax/jax/lib/xla_bridge.py:144: UserWarning: No GPU/TPU found, falling back to CPU.
warnings.warn('No GPU/TPU found, falling back to CPU.')
Traceback (most recent call last):
File "loop.py", line 11, in <module>
print(lax.while_loop(cond, body, (33, 4)))
File "/usr/local/google/home/phawkins/p/jax/jax/lax/lax_control_flow.py", line 108, in while_loop
aval_out=aval_out, cond_jaxpr=cond_jaxpr, body_jaxpr=body_jaxpr)
File "/usr/local/google/home/phawkins/p/jax/jax/core.py", line 75, in bind
return self.impl(*args, **kwargs)
File "/usr/local/google/home/phawkins/p/jax/jax/interpreters/xla.py", line 51, in apply_primitive
compiled_fun = xla_primitive_callable(prim, *abstract_args, **kwargs)
File "/usr/local/google/home/phawkins/p/jax/jax/util.py", line 174, in memoized_fun
ans = cache[key] = fun(*args, **kwargs)
File "/usr/local/google/home/phawkins/p/jax/jax/interpreters/xla.py", line 57, in xla_primitive_callable
built_c = primitive_computation(prim, *shapes, **kwargs)
File "/usr/local/google/home/phawkins/p/jax/jax/util.py", line 174, in memoized_fun
ans = cache[key] = fun(*args, **kwargs)
File "/usr/local/google/home/phawkins/p/jax/jax/interpreters/xla.py", line 74, in primitive_computation
raise e
File "/usr/local/google/home/phawkins/p/jax/jax/interpreters/xla.py", line 70, in primitive_computation
return c.Build()
File "/usr/local/google/home/phawkins/p/jax/jax/lib/xla_bridge.py", line 283, in Build
*args, **kwargs)
File "/usr/local/google/home/phawkins/.pyenv/versions/3.7.3/lib/python3.7/site-packages/jaxlib/xla_client.py", line 875, in Build
return Computation(self._builder.Build(), backend=backend)
RuntimeError: Invalid argument: The parameter of condition and body, the result of the body, and init must all have the same shape; got Condition: (parameter: ((s32[], s32[]), (s32[]), (s32[]))) -> pred[]; body: (parameter: ((s32[], s32[]), (s32[]), (s32[]))) -> (((), s32[]), (s32[]), (s32[])); init: ((s32[], s32[]), (s32[]), (s32[]))..:

# Fixed code
def xla_callable(fun, *abstract_args):
    pvals = [pe.PartialVal((aval, core.unit)) for aval in abstract_args]
    with core.new_master(pe.JaxprTrace, True) as master:
        jaxpr, (pval, consts, env) = pe.trace_to_subjaxpr(
            fun, master, False
        ).call_wrapped(pvals)
        assert not env  # no subtraces here (though cond might eventually need them)
        compiled, result_shape = compile_jaxpr(jaxpr, consts, *abstract_args)
        del master, consts, jaxpr, env
    handle_result = result_handler(result_shape)
    return partial(execute_compiled, compiled, pval, handle_result)