<p dir="auto">One of the biggest annoyances with using the fantastic <code class="notranslate">host_callback.id_tap</code> as a logging solution is its behaviour in a <code class="notranslate">jax.remat</code> (making a call to <code class="notranslate">id_tap</code> for every <code class="notranslate">remat</code>). When you have a codebase with multiple layers of nested <code class="notranslate">remat</code>s, you get spammed with calls to <code class="notranslate">host_callback.id_tap</code>, which both increases the chances of a slowdown due to Python callback overhead, but also makes it hard to efficiently get rid of spurious duplicated logs coming from <code class="notranslate">remat</code>s.</p>
<p dir="auto">Is there any way of implementing an option to <code class="notranslate">id_tap</code> to avoid calling it multiple times in a <code class="notranslate">remat</code>, or some other workaround?</p>
<p dir="auto"><a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/mattjj/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/mattjj">@mattjj</a> , <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/gnecula/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/gnecula">@gnecula</a></p>
<p dir="auto">To support modern packaging standards (aka <a href="https://peps.python.org/pep-0518/" rel="nofollow">PEP 518</a>) for wheel building, specify the build requirements and backend in <a href="https://packaging.python.org/en/latest/glossary/#term-pyproject.toml" rel="nofollow">pyproject.toml</a>.</p>
<p dir="auto">I've tested following the instructions in the <a href="https://setuptools.pypa.io/en/latest/userguide/quickstart.html#basic-use" rel="nofollow">setuptools documentation</a> (adding the <code class="notranslate">build-system</code> table), and installation from sdist works perfectly.</p>
<p dir="auto">This will fix installation when there are no wheels provided and <a href="https://pypi.org/project/setuptools/" rel="nofollow">setuptools</a> is not yet installed.</p>