Turning on fine-grained PJRT client logging

gistfile1.txt

skyewm@t1v-n-0c604497-w-0:~$ TF_CPP_MIN_LOG_LEVEL=0 TF_CPP_VMODULE=pjrt_stream_executor_client=1 python3 -c "import jax; print(jax.numpy.add(1,1))"
2022-08-31 18:03:35.722123: I external/org_tensorflow/tensorflow/compiler/xla/service/service.cc:173] XLA service 0x2c13080 initialized for platform Interpreter (this does not guarantee that XLA will be used). Devices:
2022-08-31 18:03:35.722165: I external/org_tensorflow/tensorflow/compiler/xla/service/service.cc:181]   StreamExecutor device (0): Interpreter, <undefined>
2022-08-31 18:03:35.785537: I external/org_tensorflow/tensorflow/compiler/xla/pjrt/tfrt_cpu_pjrt_client.cc:181] TfrtCpuClient created.
2022-08-31 18:03:35.786313: I external/org_tensorflow/tensorflow/core/tpu/tpu_initializer_helper.cc:262] Libtpu path is: /home/skyewm/.local/lib/python3.8/site-packages/libtpu/libtpu.so
2022-08-31 18:03:41.758624: I external/org_tensorflow/tensorflow/compiler/xla/service/service.cc:173] XLA service 0x536d500 initialized for platform TPU (this does not guarantee that XLA will be used). Devices:
2022-08-31 18:03:41.758675: I external/org_tensorflow/tensorflow/compiler/xla/service/service.cc:181]   StreamExecutor device (0): TPU, 2a886c8
2022-08-31 18:03:41.758680: I external/org_tensorflow/tensorflow/compiler/xla/service/service.cc:181]   StreamExecutor device (1): TPU, 2a886c8
2022-08-31 18:03:41.758685: I external/org_tensorflow/tensorflow/compiler/xla/service/service.cc:181]   StreamExecutor device (2): TPU, 2a886c8
2022-08-31 18:03:41.758689: I external/org_tensorflow/tensorflow/compiler/xla/service/service.cc:181]   StreamExecutor device (3): TPU, 2a886c8
2022-08-31 18:03:41.826443: I external/org_tensorflow/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.cc:2449] PjRtStreamExecutorClient::Compile
2022-08-31 18:03:41.875592: I external/org_tensorflow/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.cc:698] PjRtStreamExecutorClient::BufferFromHostBuffer: shape: s32[] device: TPU_0(process=0,(0,0,0,0))
2022-08-31 18:03:41.875700: I external/org_tensorflow/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.cc:698] PjRtStreamExecutorClient::BufferFromHostBuffer: shape: s32[] device: TPU_0(process=0,(0,0,0,0))
2022-08-31 18:03:41.875769: I external/org_tensorflow/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.cc:2217] Executing computation jit_fn.0; num_replicas=1 num_partitions=1 num_addressable_devices=1
2022-08-31 18:03:41.875788: I external/org_tensorflow/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.cc:2116] Replica 0, partition 0 mapped to device ordinal for execution: 0
2022-08-31 18:03:41.876197: I external/org_tensorflow/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.cc:2006] Replica 0 partition 0 completed; ok=1
2022-08-31 18:03:41.876245: I external/org_tensorflow/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.cc:2284] Replicated execution complete.
2022-08-31 18:03:41.876304: I external/org_tensorflow/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.cc:1224] PjRtStreamExecutorBuffer::Delete
2022-08-31 18:03:41.876329: I external/org_tensorflow/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.cc:1224] PjRtStreamExecutorBuffer::Delete
2022-08-31 18:03:41.876425: I external/org_tensorflow/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.cc:1314] PjRtStreamExecutorBuffer::ToLiteral
2
2022-08-31 18:03:41.877527: I external/org_tensorflow/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.cc:1224] PjRtStreamExecutorBuffer::Delete
2022-08-31 18:03:41.950722: I external/org_tensorflow/tensorflow/compiler/xla/pjrt/tfrt_cpu_pjrt_client.cc:184] TfrtCpuClient destroyed.