Issue 1:
Please:

[ X] Check for duplicate issues.

[X ] Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:

Running a colab notebook (here) based in Megvi's Yolox stopped working since 2 days ago because of an issue with the _can_use_color() from pretty_printer.py. The error is created by the exception of get_ipython().that if we print the traceback inside _can_use_color() reports get_ipython() not found.

Please note that for the last several weeks the _can_use_color() did not raise the exception of NameError for get_ipython, working like a charm. Might colab have changed something?

def _can_use_color() -> bool:
try:
# Check if we're in IPython or Colab
ipython = get_ipython() # type: ignore[name-defined]
shell = ipython.__class__.__name__
if shell == "ZMQInteractiveShell":
# Jupyter Notebook
return True
elif "colab" in str(ipython.__class__):
# Google Colab (external or internal)
return True
except NameError:
pass
# Otherwise check if we're in a terminal
return sys.stdout.isatty()

If applicable, include full error messages/tracebacks.

File "/usr/local/lib/python3.7/dist-packages/tensorflow/lite/python/util.py", line 51, in <module>
from jax import xla_computation as _xla_computation

File "/usr/local/lib/python3.7/dist-packages/jax/__init__.py", line 59, in <module>
from .core import eval_context as ensure_compile_time_eval

File "/usr/local/lib/python3.7/dist-packages/jax/core.py", line 47, in <module>
import jax._src.pretty_printer as pp

File "/usr/local/lib/python3.7/dist-packages/jax/_src/pretty_printer.py", line 61, in <module>
CAN_USE_COLOR = _can_use_color()
└ <function _can_use_color at 0x7f31629d3320>

File "/usr/local/lib/python3.7/dist-packages/jax/_src/pretty_printer.py", line 59, in _can_use_color
return sys.stdout.isatty()
│ └ <yolox.utils.logger.StreamToLoguru object at 0x7f353ea0e390>
└ <module 'sys' (built-in)>

AttributeError: 'StreamToLoguru' object has no attribute 'isatty'

The error of object has no attribute isatty is not important. It is a result of running the return sys.stdout.isatty() due to the exception of

Issue 2:
Dear JAX people, I've been struggling again with a leaked tracer.

The last time I had a leaked tracer, it took me a few days to find it. This time, it's been a few days and I still can't find the problem. I don't know how to make progress here. Debugging leaked tracers is hard because the error happens so far from the problem.

In the code below, I've tried to modify JAX so that it names the tracers to give me an idea of where the tracer is created. Unfortuantely, it doesn't seem to retain the names probably because tracers are copied around and the name is lost. Any insight on this would be really helpful.

I realize that tracers are "deep internals" as Matt put it. I realize you don't want to document them. Would it be possible to improve the debug messages produced? Maybe with a flag?

from inspect import signature
from typing import Any, Callable, Union

from functools import wraps, partial
from jax.interpreters.partial_eval import JaxprTracer, Trace
from jax import jit, grad, vjp
from jax.tree_util import tree_map

__all__ = ['set_debug_reprs', 'djit']

def _id_str(obj: Any) -> str:
id_ = hex(id(obj) & 0xfffff)[-5:]
if hasattr(obj, '_name'):
return f'{obj._name}:{id_}'
return id_

def set_debug_reprs():
def trace_repr(self: Trace):
level = f'{self.level}/{self.sublevel}'
return f'{self.__class__.__name__}({_id_str(self)}:{level})'

def jaxpr_tracer_repr(self: JaxprTracer):
trace = self._trace
trace_id = _id_str(trace)
return f'Tracer<{trace_id}::{_id_str(self)}>'

Trace.__repr__ = trace_repr
JaxprTracer.__repr__ = jaxpr_tracer_repr

def name_tracer(function_name, argument_name, tracer):
if not isinstance(tracer, JaxprTracer):
return tracer
tracer._name = argument_name
tracer._trace._name = function_name
print(tracer)
return tracer

def function_name(fun: Callable) -> str:
if hasattr(fun, '__name__'):
return fun.__name__
if hasattr(fun, 'func'):
return function_name(fun.func)
return "NAMELESS"

def name_all_tracers(fun):
s = signature(fun)
@wraps(fun)
def new_fun(*args, **kwargs):
bound_arguments = s.bind(*args, **kwargs)
for name, value in bound_arguments.arguments.items():
tree_map(partial(name_tracer, function_name(fun), name), value)
return fun(*args, **kwargs)
return new_fun

def djit(fun: Callable, *args, **kwargs):
new_fun = name_all_tracers(fun)
return jit(new_fun, *args, **kwargs)

def dgrad(fun, *args, **kwargs):
new_fun = name_all_tracers(fun)
return grad(new_fun, *args, **kwargs)

def dvjp(fun, *args, **kwargs):
new_fun = name_all_tracers(fun)
return vjp(new_fun, *args, **kwargs)

When debugging, I call set_debug_reprs() and use the shim:

from jax import *
from .debugging import djit as jit, dgrad as grad, dvjp as vjp

With this, a JaxprTracer now prints out as:

Tracer<find_fixed_point:84220::x_init:87810>

This shows the function name in which it was created find_fixed_point, the parameter name that it corresponds to x_init, and some object id codes for each in case there are duplicate invocations.

It would have been nice to also name the field in case the tracer is part of a PyTree-like object, but that's hard to get unless the PyTree-like object exposes names for its fields like in flax.struct.dataclass that was recently shown to me. (Perhaps the PyTree registry could optionally accept a third function that produces a tree of strings for naming purposes.)