# Buggy code
def parallel_callable(fun, axis_name, axis_size, *avals):
    pvals = [PartialVal((aval, core.unit)) for aval in avals]
    with core.new_master(JaxprTrace, True) as master:
        jaxpr, (pval, consts, env) = trace_to_subjaxpr(fun, master).call_wrapped(pvals)
        assert not env
        out = compile_replicated(jaxpr, axis_name, axis_size, consts, *avals)
        compiled, nrep, result_shape = out
        del master, consts, jaxpr, env
    handle_arg = partial(shard_arg, compiled._device_ordinals)
    handle_result = xla.result_handler(result_shape)
    return partial(
        execute_replicated, compiled, pval, axis_size, nrep, handle_arg, 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 parallel_callable(fun, axis_name, axis_size, *avals):
    pvals = [PartialVal((aval, core.unit)) for aval in avals]
    with core.new_master(JaxprTrace, True) as master:
        jaxpr, (pval, consts, env) = trace_to_subjaxpr(fun, master, False).call_wrapped(
            pvals
        )
        assert not env
        out = compile_replicated(jaxpr, axis_name, axis_size, consts, *avals)
        compiled, nrep, result_shape = out
        del master, consts, jaxpr, env
    handle_arg = partial(shard_arg, compiled._device_ordinals)
    handle_result = xla.result_handler(result_shape)
    return partial(
        execute_replicated, compiled, pval, axis_size, nrep, handle_arg, handle_result
    )