How to detect unused parameters passed to a jitted function?

[jjyyxx]

When I pass a large PyTree of neural network parameters to a jitted function, it may contain certain leaves that is not used anywhere inside that function. Sometimes, this indicates a bug in code. So, is there any way to detect this?

For example, haiku creates parameters for all get_parameter calls. So, params will contain weights that is eventually unused (after jit optimization of f_tr.apply)

def f(x):
    # Forget to assign the output to x
    hk.Linear(256)(x)
    return x

key = jax.random.PRNGKey(42)
x = jax.random.normal(key, [256])
f_tr = hk.transform(f)
params = jax.jit(f_tr.init)(key, x)
jax.jit(f_tr.apply)(params, None, x)

[jakevdp]

I'm not sure what you mean by "detect" here, but one way to see which parameters in your computation are unused is to print the jaxpr and look for parameters that are assigned to _:

print(jax.make_jaxpr(f_tr.apply)(params, None, x))

{ lambda ; a:f32[256] b:f32[256,256] c:f32[256]. let
    d:f32[256] = dot_general[
      dimension_numbers=(([0], [0]), ([], []))
      preferred_element_type=float32
    ] c b
    _:f32[256] = add d a
  in (c,) }

These are among the parameters that will be eliminated as dead code by the compiler before execution.

Does that answer your question?