omnistaging experiment: print partial eval stats

[mattjj]

The idea here is that we can add a useful debugging tool for omnistaging (#3370, #4038) by printing out expressions that would have been constant-folded with pre-omnistaging behavior.

import jax.numpy as jnp
from jax.lax import scan

from jax.config import config
config.enable_omnistaging()

def cumsum(xs):
  _, ys = scan(lambda c, x: (c + x, c + x + jnp.add(1, 1)), 0., xs)
  return ys

cumsum(jnp.arange(10))

from tracing <lambda> at example.py:8, this was not constant folded:

{ lambda  ; _ _.
  let c = add 1 1
      d = convert_element_type[ new_dtype=float32
                                old_dtype=int32 ] c
  in (*, *, d) }

[hawkinsp]

Closing, because I doubt many folks need this particular debugging facility any more...