Support profiling without source code changes

Currently adding profiling requires source-code changes which could be difficult in a large codebase.

Would be great if something like JAX_PROFILE=output.pb was supported to generate profiling from command line.

CC @hawkinsp