how to dump the operations of all the methods of a jitted class object

[ronghongbo]

Hello, I tried flax.linen.jit with a class object. It seems to work without any error, but then how to dump out the operations/IR of the methods in the jitted object?

Here is a simple program mlp.py:

from typing import Sequence

import numpy as np
import jax
import jax.numpy as jnp
import flax.linen as nn

class MLP(nn.Module):
  features: Sequence[int]

  @nn.compact
  def __call__(self, x):
    for feat in self.features[:-1]:
      x = nn.relu(nn.Dense(feat)(x))
    x = nn.Dense(self.features[-1])(x)
    return x

model = MLP([12, 8, 4])
batch = jnp.ones((32, 10))
variables = model.init(jax.random.PRNGKey(0), batch)
output = model.apply(variables, batch)

# Here is to jit the the model object. I am not sure how to dump the results, so I just print it.
w=nn.jit(model)
print(w)

Here is the execution results:

$ python mlp.py
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
<function wrap_method_once.<locals>.wrapped_module_method at 0x7f757c1e9e50>

The results provide an address for the wrapped method, instead of a human readable IR.

Any idea?

Thanks!
Hongbo

[ronghongbo]

Thanks to the help of Bojan Nikolic, I was able to dump the IR in this way:

import numpy as np
import jax
import jax.numpy as jnp
import flax.linen as nn
import jax.tree_util as tree_util
import jax.tools.jax_to_ir
import functools
from typing import Sequence
from functools import partial
from jax.lib import xla_client

# The original code
class MLP(nn.Module):
  features: Sequence[int]

  @nn.compact
  def __call__(self, x):
    for feat in self.features[:-1]:
      x = nn.relu(nn.Dense(feat)(x))
    x = nn.Dense(self.features[-1])(x)
    return x

model = MLP([12, 8, 4])
batch = jnp.ones((32, 10))
variables = model.init(jax.random.PRNGKey(0), batch)
output = model.apply(variables, batch)

# New code to dump IR
f=functools.partial(model.apply, variables)
z=jax.xla_computation(f)(batch)
u=xla_client._xla.mlir.xla_computation_to_mlir_module(z)
print(u)