Manually inspecting jitted code

[vlad17]

Is there a way that I can inspect jax.jitted code? E.g., seeing the generated XLA HLO IR or LLVM IR for cpu targets.

This came up while I was writing a cdist implementation in jax (if you're curious, for vectorizing VLAD, Colab).

The cdist implementation is as follows:

import jax
import jax.numpy as jnp

def sqeuclidean(a, b):
  return jnp.sum((a - b) ** 2)

# *two* vmap calls necessary to allow for pairwise distances between two
# collections of row vectors of different sizes
cdist = jax.vmap(jax.vmap(sqeuclidean, in_axes=(None, 0)), in_axes=(0, None))

I am moderately confident that jax.jit(sqeuclidean) would result in an appropriately-fused implementation which doesn't generate a large intermediate tensor (a-b)**2. However, I'm less confident that this would be the case for jax.jit(cdist). Either way, being able to inspect IR would let me confirm. Is there a way I could do that?

[jakevdp]

You can output non-optimized or optimized HLO using a function like the one in this pull request: #7719 (only a few lines, so easily copy-paste-able)

It hasn't been merged because we could not come to a consensus on the best API for this functionality, but more complete work on this kind of thing is being tracked in #7733

[vlad17]

Great! I'm excited for that functionality to come to main.