<p dir="auto">Please:</p>
<ul dir="auto">
<li>[ X] Check for duplicate issues.</li>
<li>[X ] Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:</li>
</ul>
<p dir="auto">Running a colab notebook (<a href="https://github.com/Kazuhito00/YOLOX-Colaboratory-Training-Sample/blob/main/YOLOX_Colaboratory_Training_Sample.ipynb">here</a>) 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.</p>
<p dir="auto">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?</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="def _can_use_color() -&gt; bool:
  try:
    # Check if we're in IPython or Colab
    ipython = get_ipython()  # type: ignore[name-defined]
    shell = ipython.__class__.__name__
    if shell == &quot;ZMQInteractiveShell&quot;:
      # Jupyter Notebook
      return True
    elif &quot;colab&quot; 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()"><pre class="notranslate"><code class="notranslate">def _can_use_color() -&gt; 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()
</code></pre></div>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox"> If applicable, include full error messages/tracebacks.</li>
</ul>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="  File &quot;/usr/local/lib/python3.7/dist-packages/tensorflow/lite/python/util.py&quot;, line 51, in &lt;module&gt;
    from jax import xla_computation as _xla_computation

  File &quot;/usr/local/lib/python3.7/dist-packages/jax/__init__.py&quot;, line 59, in &lt;module&gt;
    from .core import eval_context as ensure_compile_time_eval

  File &quot;/usr/local/lib/python3.7/dist-packages/jax/core.py&quot;, line 47, in &lt;module&gt;
    import jax._src.pretty_printer as pp

  File &quot;/usr/local/lib/python3.7/dist-packages/jax/_src/pretty_printer.py&quot;, line 61, in &lt;module&gt;
    CAN_USE_COLOR = _can_use_color()
                    └ &lt;function _can_use_color at 0x7f31629d3320&gt;

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

AttributeError: 'StreamToLoguru' object has no attribute 'isatty'"><pre class="notranslate"><code class="notranslate">  File "/usr/local/lib/python3.7/dist-packages/tensorflow/lite/python/util.py", line 51, in &lt;module&gt;
    from jax import xla_computation as _xla_computation

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

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

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

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

AttributeError: 'StreamToLoguru' object has no attribute 'isatty'
</code></pre></div>
<p dir="auto">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</p>
<p dir="auto">Dear JAX people, I've been struggling again with a leaked tracer.</p>
<p dir="auto"><a href="https://github.com/google/jax/issues/2912" data-hovercard-type="issue" data-hovercard-url="/google/jax/issues/2912/hovercard">The last time I had a leaked tracer</a>, 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.</p>
<p dir="auto">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.</p>
<p dir="auto">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?</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="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) -&gt; str:
    id_ = hex(id(obj) &amp; 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&lt;{trace_id}::{_id_str(self)}&gt;'

    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) -&gt; str:
    if hasattr(fun, '__name__'):
        return fun.__name__
    if hasattr(fun, 'func'):
        return function_name(fun.func)
    return &quot;NAMELESS&quot;


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)"><pre class="notranslate"><span class="pl-k">from</span> <span class="pl-s1">inspect</span> <span class="pl-k">import</span> <span class="pl-s1">signature</span>
<span class="pl-k">from</span> <span class="pl-s1">typing</span> <span class="pl-k">import</span> <span class="pl-v">Any</span>, <span class="pl-v">Callable</span>, <span class="pl-v">Union</span>

<span class="pl-k">from</span> <span class="pl-s1">functools</span> <span class="pl-k">import</span> <span class="pl-s1">wraps</span>, <span class="pl-s1">partial</span>
<span class="pl-k">from</span> <span class="pl-s1">jax</span>.<span class="pl-s1">interpreters</span>.<span class="pl-s1">partial_eval</span> <span class="pl-k">import</span> <span class="pl-v">JaxprTracer</span>, <span class="pl-v">Trace</span>
<span class="pl-k">from</span> <span class="pl-s1">jax</span> <span class="pl-k">import</span> <span class="pl-s1">jit</span>, <span class="pl-s1">grad</span>, <span class="pl-s1">vjp</span>
<span class="pl-k">from</span> <span class="pl-s1">jax</span>.<span class="pl-s1">tree_util</span> <span class="pl-k">import</span> <span class="pl-s1">tree_map</span>

<span class="pl-s1">__all__</span> <span class="pl-c1">=</span> [<span class="pl-s">'set_debug_reprs'</span>, <span class="pl-s">'djit'</span>]


<span class="pl-k">def</span> <span class="pl-en">_id_str</span>(<span class="pl-s1">obj</span>: <span class="pl-v">Any</span>) <span class="pl-c1">-&gt;</span> <span class="pl-s1">str</span>:
    <span class="pl-s1">id_</span> <span class="pl-c1">=</span> <span class="pl-en">hex</span>(<span class="pl-en">id</span>(<span class="pl-s1">obj</span>) <span class="pl-c1">&amp;</span> <span class="pl-c1">0xfffff</span>)[<span class="pl-c1">-</span><span class="pl-c1">5</span>:]
    <span class="pl-k">if</span> <span class="pl-en">hasattr</span>(<span class="pl-s1">obj</span>, <span class="pl-s">'_name'</span>):
        <span class="pl-k">return</span> <span class="pl-s">f'<span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">obj</span>.<span class="pl-s1">_name</span><span class="pl-kos">}</span></span>:<span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">id_</span><span class="pl-kos">}</span></span>'</span>
    <span class="pl-k">return</span> <span class="pl-s1">id_</span>


<span class="pl-k">def</span> <span class="pl-en">set_debug_reprs</span>():
    <span class="pl-k">def</span> <span class="pl-en">trace_repr</span>(<span class="pl-s1">self</span>: <span class="pl-v">Trace</span>):
        <span class="pl-s1">level</span> <span class="pl-c1">=</span> <span class="pl-s">f'<span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">self</span>.<span class="pl-s1">level</span><span class="pl-kos">}</span></span>/<span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">self</span>.<span class="pl-s1">sublevel</span><span class="pl-kos">}</span></span>'</span>
        <span class="pl-k">return</span> <span class="pl-s">f'<span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">self</span>.<span class="pl-s1">__class__</span>.<span class="pl-s1">__name__</span><span class="pl-kos">}</span></span>(<span class="pl-s1"><span class="pl-kos">{</span><span class="pl-en">_id_str</span>(<span class="pl-s1">self</span>)<span class="pl-kos">}</span></span>:<span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">level</span><span class="pl-kos">}</span></span>)'</span>

    <span class="pl-k">def</span> <span class="pl-en">jaxpr_tracer_repr</span>(<span class="pl-s1">self</span>: <span class="pl-v">JaxprTracer</span>):
        <span class="pl-s1">trace</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-s1">_trace</span>
        <span class="pl-s1">trace_id</span> <span class="pl-c1">=</span> <span class="pl-en">_id_str</span>(<span class="pl-s1">trace</span>)
        <span class="pl-k">return</span> <span class="pl-s">f'Tracer&lt;<span class="pl-s1"><span class="pl-kos">{</span><span class="pl-s1">trace_id</span><span class="pl-kos">}</span></span>::<span class="pl-s1"><span class="pl-kos">{</span><span class="pl-en">_id_str</span>(<span class="pl-s1">self</span>)<span class="pl-kos">}</span></span>&gt;'</span>

    <span class="pl-v">Trace</span>.<span class="pl-s1">__repr__</span> <span class="pl-c1">=</span> <span class="pl-s1">trace_repr</span>
    <span class="pl-v">JaxprTracer</span>.<span class="pl-s1">__repr__</span> <span class="pl-c1">=</span> <span class="pl-s1">jaxpr_tracer_repr</span>


<span class="pl-k">def</span> <span class="pl-en">name_tracer</span>(<span class="pl-s1">function_name</span>, <span class="pl-s1">argument_name</span>, <span class="pl-s1">tracer</span>):
    <span class="pl-k">if</span> <span class="pl-c1">not</span> <span class="pl-en">isinstance</span>(<span class="pl-s1">tracer</span>, <span class="pl-v">JaxprTracer</span>):
        <span class="pl-k">return</span> <span class="pl-s1">tracer</span>
    <span class="pl-s1">tracer</span>.<span class="pl-s1">_name</span> <span class="pl-c1">=</span> <span class="pl-s1">argument_name</span>
    <span class="pl-s1">tracer</span>.<span class="pl-s1">_trace</span>.<span class="pl-s1">_name</span> <span class="pl-c1">=</span> <span class="pl-s1">function_name</span>
    <span class="pl-en">print</span>(<span class="pl-s1">tracer</span>)
    <span class="pl-k">return</span> <span class="pl-s1">tracer</span>


<span class="pl-k">def</span> <span class="pl-en">function_name</span>(<span class="pl-s1">fun</span>: <span class="pl-v">Callable</span>) <span class="pl-c1">-&gt;</span> <span class="pl-s1">str</span>:
    <span class="pl-k">if</span> <span class="pl-en">hasattr</span>(<span class="pl-s1">fun</span>, <span class="pl-s">'__name__'</span>):
        <span class="pl-k">return</span> <span class="pl-s1">fun</span>.<span class="pl-s1">__name__</span>
    <span class="pl-k">if</span> <span class="pl-en">hasattr</span>(<span class="pl-s1">fun</span>, <span class="pl-s">'func'</span>):
        <span class="pl-k">return</span> <span class="pl-en">function_name</span>(<span class="pl-s1">fun</span>.<span class="pl-s1">func</span>)
    <span class="pl-k">return</span> <span class="pl-s">"NAMELESS"</span>


<span class="pl-k">def</span> <span class="pl-en">name_all_tracers</span>(<span class="pl-s1">fun</span>):
    <span class="pl-s1">s</span> <span class="pl-c1">=</span> <span class="pl-en">signature</span>(<span class="pl-s1">fun</span>)
    <span class="pl-en">@<span class="pl-en">wraps</span>(<span class="pl-s1">fun</span>)</span>
    <span class="pl-k">def</span> <span class="pl-en">new_fun</span>(<span class="pl-c1">*</span><span class="pl-s1">args</span>, <span class="pl-c1">**</span><span class="pl-s1">kwargs</span>):
        <span class="pl-s1">bound_arguments</span> <span class="pl-c1">=</span> <span class="pl-s1">s</span>.<span class="pl-en">bind</span>(<span class="pl-c1">*</span><span class="pl-s1">args</span>, <span class="pl-c1">**</span><span class="pl-s1">kwargs</span>)
        <span class="pl-k">for</span> <span class="pl-s1">name</span>, <span class="pl-s1">value</span> <span class="pl-c1">in</span> <span class="pl-s1">bound_arguments</span>.<span class="pl-s1">arguments</span>.<span class="pl-en">items</span>():
            <span class="pl-en">tree_map</span>(<span class="pl-en">partial</span>(<span class="pl-s1">name_tracer</span>, <span class="pl-en">function_name</span>(<span class="pl-s1">fun</span>), <span class="pl-s1">name</span>), <span class="pl-s1">value</span>)
        <span class="pl-k">return</span> <span class="pl-en">fun</span>(<span class="pl-c1">*</span><span class="pl-s1">args</span>, <span class="pl-c1">**</span><span class="pl-s1">kwargs</span>)
    <span class="pl-k">return</span> <span class="pl-s1">new_fun</span>


<span class="pl-k">def</span> <span class="pl-en">djit</span>(<span class="pl-s1">fun</span>: <span class="pl-v">Callable</span>, <span class="pl-c1">*</span><span class="pl-s1">args</span>, <span class="pl-c1">**</span><span class="pl-s1">kwargs</span>):
    <span class="pl-s1">new_fun</span> <span class="pl-c1">=</span> <span class="pl-en">name_all_tracers</span>(<span class="pl-s1">fun</span>)
    <span class="pl-k">return</span> <span class="pl-en">jit</span>(<span class="pl-s1">new_fun</span>, <span class="pl-c1">*</span><span class="pl-s1">args</span>, <span class="pl-c1">**</span><span class="pl-s1">kwargs</span>)


<span class="pl-k">def</span> <span class="pl-en">dgrad</span>(<span class="pl-s1">fun</span>, <span class="pl-c1">*</span><span class="pl-s1">args</span>, <span class="pl-c1">**</span><span class="pl-s1">kwargs</span>):
    <span class="pl-s1">new_fun</span> <span class="pl-c1">=</span> <span class="pl-en">name_all_tracers</span>(<span class="pl-s1">fun</span>)
    <span class="pl-k">return</span> <span class="pl-en">grad</span>(<span class="pl-s1">new_fun</span>, <span class="pl-c1">*</span><span class="pl-s1">args</span>, <span class="pl-c1">**</span><span class="pl-s1">kwargs</span>)


<span class="pl-k">def</span> <span class="pl-en">dvjp</span>(<span class="pl-s1">fun</span>, <span class="pl-c1">*</span><span class="pl-s1">args</span>, <span class="pl-c1">**</span><span class="pl-s1">kwargs</span>):
    <span class="pl-s1">new_fun</span> <span class="pl-c1">=</span> <span class="pl-en">name_all_tracers</span>(<span class="pl-s1">fun</span>)
    <span class="pl-k">return</span> <span class="pl-en">vjp</span>(<span class="pl-s1">new_fun</span>, <span class="pl-c1">*</span><span class="pl-s1">args</span>, <span class="pl-c1">**</span><span class="pl-s1">kwargs</span>)</pre></div>
<p dir="auto">When debugging, I call <code class="notranslate">set_debug_reprs()</code> and use the shim:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="from jax import *
from .debugging import djit as jit, dgrad as grad, dvjp as vjp"><pre class="notranslate"><span class="pl-k">from</span> <span class="pl-s1">jax</span> <span class="pl-k">import</span> <span class="pl-c1">*</span>
<span class="pl-k">from</span> .<span class="pl-s1">debugging</span> <span class="pl-k">import</span> <span class="pl-s1">djit</span> <span class="pl-k">as</span> <span class="pl-s1">jit</span>, <span class="pl-s1">dgrad</span> <span class="pl-k">as</span> <span class="pl-s1">grad</span>, <span class="pl-s1">dvjp</span> <span class="pl-k">as</span> <span class="pl-s1">vjp</span></pre></div>
<p dir="auto">With this, a <code class="notranslate">JaxprTracer</code> now prints out as:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="Tracer&lt;find_fixed_point:84220::x_init:87810&gt;"><pre class="notranslate"><code class="notranslate">Tracer&lt;find_fixed_point:84220::x_init:87810&gt;
</code></pre></div>
<p dir="auto">This shows the function name in which it was created <code class="notranslate">find_fixed_point</code>, the parameter name that it corresponds to <code class="notranslate">x_init</code>, and some object id codes for each in case there are duplicate invocations.</p>
<p dir="auto">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 <code class="notranslate">flax.struct.dataclass</code> 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.)</p>