Is there any way to get the computational graph for a function with Python while control flow?

[ToshiyukiBandai]

Hello everyone,

I am trying to visualize the computational graph for a function with Python while control flow (I followed a guide by https://bnikolic.co.uk/blog/python/jax/2022/02/22/jax-outputgraph-rev.html). For example, for a function

def test_func(x):
    while x > 0.1:
        x = x * 0.5
    return x

I tried to use z = jax.xla_computation(test_func)(3.0), but it did not work because of the value dependent loop. The error message was

ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: traced array with shape bool[].
The problem arose with the `bool` function. 

Is there any way to produce the HLO intermediate representation of such functions? For example, can I explicitly unroll the function and feed the unrolled function into jax.xla_computation?

[jakevdp]

Data-dependent python while loops can only be executed eagerly, so I don't know of any easy way to lower the computations to an HLO IR.

[ToshiyukiBandai]

Thank you so much for your quick response. I asked this question because this video around 22:20 (https://www.youtube.com/watch?v=MX1RJELWONc) showed the computational graph for a function with a Python while loop (they used "grad_graph" function), so I thought there might be a way to do that...
But, thank you for your answer.

[jakevdp]

@mattjj might remember what the grad_graph function does.