<p dir="auto">In v0.2.11 jax added the nice optional installation targets <code class="notranslate">cpu</code>, <code class="notranslate">cudaXXX</code> and <code class="notranslate">tpu</code> which install jaxlib for the desired platform. Currently, the setup script guarantees that the installed version of jaxlib is compatible by constraining the version number from below to the earliest known compatible version. However, no similar constraint is put for a maximum compatible version number. Given that jaxlib does not appear to necessarily maintain backward compatibility between updates, this is an issue for installing older versions of jax (as these will always pull the latest available jaxlib version, leading to crashes due to incompatibility).</p>
<p dir="auto">Jax installation should always result in a working setup, i.e., only the latest compatible version of jaxlib should be installed when installing jax. This either requires version constraints in jax's setup.py or backward compatibility guarantees in jaxlib (or a combination of both).</p>
<p dir="auto">Please:</p>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Check for duplicate requests.</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Describe your goal, and if possible provide a code snippet with a motivating example.</li>
</ul>
<p dir="auto">Consider the following code:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="@dataclasses.dataclass
class Config:
  some_value: int
  # many other values omitted

@functools.partial(jax.jit, static_argnums=1)
def do_stuff(data, config):
  # do some stuff

def evaluate():
  config = Config(suitable_value_for_evaluation)
  data = summon_up_evaluation_data()
  do_stuff(data, config)"><pre class="notranslate"><span class="pl-en">@<span class="pl-s1">dataclasses</span>.<span class="pl-s1">dataclass</span></span>
<span class="pl-k">class</span> <span class="pl-v">Config</span>:
  <span class="pl-s1">some_value</span>: <span class="pl-s1">int</span>
  <span class="pl-c"># many other values omitted</span>

<span class="pl-en">@<span class="pl-s1">functools</span>.<span class="pl-en">partial</span>(<span class="pl-s1">jax</span>.<span class="pl-s1">jit</span>, <span class="pl-s1">static_argnums</span><span class="pl-c1">=</span><span class="pl-c1">1</span>)</span>
<span class="pl-k">def</span> <span class="pl-en">do_stuff</span>(<span class="pl-s1">data</span>, <span class="pl-s1">config</span>):
  <span class="pl-c"># do some stuff</span>

<span class="pl-k">def</span> <span class="pl-en">evaluate</span>():
  <span class="pl-s1">config</span> <span class="pl-c1">=</span> <span class="pl-v">Config</span>(<span class="pl-s1">suitable_value_for_evaluation</span>)
  <span class="pl-s1">data</span> <span class="pl-c1">=</span> <span class="pl-en">summon_up_evaluation_data</span>()
  <span class="pl-en">do_stuff</span>(<span class="pl-s1">data</span>, <span class="pl-s1">config</span>)</pre></div>
<p dir="auto">This will unexpectedly (to me at least) retrace/recompile <code class="notranslate">do_stuff</code> every time <code class="notranslate">evaluate</code> is called. The reason is because the <code class="notranslate">dataclass</code> derived <code class="notranslate">Config</code> does not have a <code class="notranslate">__hash__</code> method by default. (Using <code class="notranslate">@dataclasses.dataclass(frozen=True)</code> does yield a <code class="notranslate">Config</code> class with a hash, and then behaviour is as expected.)</p>
<p dir="auto">It would be helpful to warn, or even error, if an unhashable object is passed to <code class="notranslate">jax.jit</code>, to make this sort of thing easier to track down. (It took me a long while before I thought to check whether dataclasses were hashable ...)</p>
<p dir="auto">It looks as though <code class="notranslate">wrap_hashably</code> being too lenient is the root cause: <a href="https://github.com/google/jax/blob/master/jax/api_util.py#L130">https://github.com/google/jax/blob/master/jax/api_util.py#L130</a> (Thanks <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/tomhennigan/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/tomhennigan">@tomhennigan</a> ). It could be better to fail in this function if the unhashable value isn't one of the types that this function intends to wrap, as it seems that the retrace/recompile behaviour would never be the one that user desired.</p>