Migrate away from absl logging.

I would suggest a small change. It seems like Jax’s logging facilities rely on Abseil like tensorflow, which is an interesting google product. But, is it really necessary to use though? Python already provides a logging mechanism, would it be easy to migrate to that? It seems to be used in

1. jaxlib/tpu_client.py
2. jaxlib/xla_client.py
3. jax/experimental/host_callback.py
4. jax/interpreters/pxla.py
5. jax/interpreters/sharded_jit.py
6. jax/interpreters/xla.py
7. jax/lib/xla_bridge.py

Using abseil can confuse users who are trying to silence warnings and assume jax is using python’s warning systems. For example, the famous

 WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

[nicholasjng]

Ideally we could still use absl optionally but hide that behind a Python logging interface. What do you think of the following?

1. Introduce a `jax.config.jax_use_absl_logging` flag, off by default.

2. Set up a module (say, `jax.logging`) that offers the interface of Python's `logging`. Have this forward to Python `logging` by default, or delegate to `absl.logging` where applicable if the config flag says to do that instead.

3. Change jax internals to use `jax.logging` across the codebase. This will require e.g. moving away from calls to absl's `vlog` and instead directly calling `jax.logging.debug` etc.

I made an attempt at implementing this solution in #10968, and would be happy about some feedback/guidance 😃

[froystig]

Ideally we could still use absl optionally but hide that behind a Python logging interface. What do you think of the following?

1. Introduce a jax.config.jax_use_absl_logging flag, off by default.
2. Set up a module (say, jax.logging) that offers the interface of Python’s logging. Have this forward to Python logging by default, or delegate to absl.logging where applicable if the config flag says to do that instead.
3. Change jax internals to use jax.logging across the codebase. This will require e.g. moving away from calls to absl’s vlog and instead directly calling jax.logging.debug etc.