jax-profiler

Memory profiler for JAX

Usage

Create profiler logs

from jaxprof import JaxProfiler

profiler = JaxProfiler()
def some_jax_code():
    ...
    profiler.capture()
    ...

or run it in the background

profiler.capture_in_background()

Generate plots from profiler logs

python jaxprof.py --help

Name		Name	Last commit message	Last commit date
Latest commit History 3 Commits
.gitignore		.gitignore
LICENSE		LICENSE
README.md		README.md
jaxprof.py		jaxprof.py