jax

1e25c44d - [mosaic_gpu] Only `jit` function to profile with cupti if it not already `jit`ted.