Possible best practice check to check retracing and unexpected slow down?

[qsh-zh]

This is a general question.

I found my jax jitted code is unexpectedly slow and memory keeps increasing. I suspect one part of jitted code is tracing multiple times. Are there any convenient tools to check retracing, and profiling?

I know I can put print in jitted function, and check how many times it prints. But do we have better approaches, I found in TF, that users can use traced_function.pretty_printed_concrete_signatures() to check. And tf.print can run every time user call the function.

[mattjj]

Thanks for the question!

You can use the jax_log_compiles option. You can set it in any of these ways:

1. Use the context manager with jax.log_compiles(True): ...
2. Use the config option with jax.config.update('jax_log_compiles', True)
3. Use a shell environment variable JAX_LOG_COMPILES=1 (or any truthy value)
4. Use absl together with the absl command-line flag --jax_log_compiles=1

We'd like to improve the tooling around this further, in particular to provide some explanation as to why a function is getting recompiled. At the moment the information is pretty minimal.

[patrick-kidger]

With respect to handling the "why": given some problematic function

@jax.jit
def big_function_mysteriously_being_recompiled(arg1, arg2, ...):
    ...

big_function_mysteriously_being_recompiled(arg1, arg2, ...)

then I define a helper function:

@functools.partial(jax.jit, static_argnum=1)
def why(arg, name):
    print(name)

and change the callsite to:

why(arg1, "arg1")
why(arg2, "arg2")
...
big_function_mysteriously_being_recompiled(arg1, arg2, ...)

which then makes clear which argument is changing.