<p dir="auto">Hi folks,<br>
Wondering if someone can please help with this import error on a fresh install.</p>
<p dir="auto">Background: Recently started a new conda environment with a couple of basics like jupyter, numpy, scipy. I'm using MacOS High Sierra 10.13.4, and python3.7</p>
<p dir="auto">The specific error is:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="---------------------------------------------------------------------------
ImportError                               Traceback (most recent call last)
&lt;ipython-input-1-77da20ac745a&gt; in &lt;module&gt;
----&gt; 1 from jax import vmap # for auto-vectorizing functions
      2 from functools import partial # for use with vmap
      3 from jax import jit # for compiling functions for speedup
      4 from jax import random # stax initialization uses jax.random
      5 from jax.experimental import stax # neural network library"><pre class="notranslate"><code class="notranslate">---------------------------------------------------------------------------
ImportError                               Traceback (most recent call last)
&lt;ipython-input-1-77da20ac745a&gt; in &lt;module&gt;
----&gt; 1 from jax import vmap # for auto-vectorizing functions
      2 from functools import partial # for use with vmap
      3 from jax import jit # for compiling functions for speedup
      4 from jax import random # stax initialization uses jax.random
      5 from jax.experimental import stax # neural network library
</code></pre></div>
<p dir="auto">...</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="~/miniconda3/envs/lew_jax/lib/python3.7/site-packages/jaxlib/xla_client.py in &lt;module&gt;
     36 # and TensorFlow may fail with duplicate protocol buffer message definitions.
     37 
---&gt; 38 from . import xla_extension as _xla
     39 from .xla_extension import ops
     40 

ImportError: dlopen(/Users/lmar3213/miniconda3/envs/lew_jax/lib/python3.7/site-packages/jaxlib/xla_extension.so, 2): Symbol not found: ____chkstk_darwin
  Referenced from: /Users/lmar3213/miniconda3/envs/lew_jax/lib/python3.7/site-packages/jaxlib/xla_extension.so (which was built for Mac OS X 10.15)
  Expected in: /usr/lib/libSystem.B.dylib
 in /Users/lmar3213/miniconda3/envs/lew_jax/lib/python3.7/site-packages/jaxlib/xla_extension.so"><pre class="notranslate"><code class="notranslate">~/miniconda3/envs/lew_jax/lib/python3.7/site-packages/jaxlib/xla_client.py in &lt;module&gt;
     36 # and TensorFlow may fail with duplicate protocol buffer message definitions.
     37 
---&gt; 38 from . import xla_extension as _xla
     39 from .xla_extension import ops
     40 

ImportError: dlopen(/Users/lmar3213/miniconda3/envs/lew_jax/lib/python3.7/site-packages/jaxlib/xla_extension.so, 2): Symbol not found: ____chkstk_darwin
  Referenced from: /Users/lmar3213/miniconda3/envs/lew_jax/lib/python3.7/site-packages/jaxlib/xla_extension.so (which was built for Mac OS X 10.15)
  Expected in: /usr/lib/libSystem.B.dylib
 in /Users/lmar3213/miniconda3/envs/lew_jax/lib/python3.7/site-packages/jaxlib/xla_extension.so
</code></pre></div>
<p dir="auto">To install, I ran as per the installation instructions i.e. <code class="notranslate">pip install --upgrade pip</code> which was already up to date, and then <code class="notranslate">pip install --upgrade jax jaxlib</code> which installed happily.</p>
<p dir="auto">The only clue I have as to understanding the error is that it mentions <code class="notranslate">xla_extension.so</code> was built for mac OS X 10.15 but Im on 10.13. Any help is appreciated! Thankyou!</p>
<h2 dir="auto">Steps to reproduce:</h2>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import functools

from absl import app
import jax
from jax.experimental import maps
from jax.experimental.pjit import PartitionSpec
from jax.experimental.pjit import pjit
import numpy as np


jax.config.update('jax_parallel_functions_output_gda', True)


@functools.partial(
    pjit,
    in_axis_resources=(),
    out_axis_resources=PartitionSpec('x'),
)
def f():
  return jax.numpy.zeros([32, 10])


def main(argv):
  del argv  # Unused
  with maps.Mesh(np.asarray(jax.devices()), axis_names=['x']):
    print(f().shape)
    print(jax.eval_shape(f))

if __name__ == '__main__':
  app.run(main)"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">functools</span>

<span class="pl-k">from</span> <span class="pl-s1">absl</span> <span class="pl-k">import</span> <span class="pl-s1">app</span>
<span class="pl-k">import</span> <span class="pl-s1">jax</span>
<span class="pl-k">from</span> <span class="pl-s1">jax</span>.<span class="pl-s1">experimental</span> <span class="pl-k">import</span> <span class="pl-s1">maps</span>
<span class="pl-k">from</span> <span class="pl-s1">jax</span>.<span class="pl-s1">experimental</span>.<span class="pl-s1">pjit</span> <span class="pl-k">import</span> <span class="pl-v">PartitionSpec</span>
<span class="pl-k">from</span> <span class="pl-s1">jax</span>.<span class="pl-s1">experimental</span>.<span class="pl-s1">pjit</span> <span class="pl-k">import</span> <span class="pl-s1">pjit</span>
<span class="pl-k">import</span> <span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">np</span>


<span class="pl-s1">jax</span>.<span class="pl-s1">config</span>.<span class="pl-en">update</span>(<span class="pl-s">'jax_parallel_functions_output_gda'</span>, <span class="pl-c1">True</span>)


<span class="pl-en">@<span class="pl-s1">functools</span>.<span class="pl-en">partial</span>(</span>
<span class="pl-en">    <span class="pl-s1">pjit</span>,</span>
<span class="pl-en">    <span class="pl-s1">in_axis_resources</span><span class="pl-c1">=</span>(),</span>
<span class="pl-en">    <span class="pl-s1">out_axis_resources</span><span class="pl-c1">=</span><span class="pl-v">PartitionSpec</span>(<span class="pl-s">'x'</span>),</span>
<span class="pl-en">)</span>
<span class="pl-k">def</span> <span class="pl-en">f</span>():
  <span class="pl-k">return</span> <span class="pl-s1">jax</span>.<span class="pl-s1">numpy</span>.<span class="pl-en">zeros</span>([<span class="pl-c1">32</span>, <span class="pl-c1">10</span>])


<span class="pl-k">def</span> <span class="pl-en">main</span>(<span class="pl-s1">argv</span>):
  <span class="pl-k">del</span> <span class="pl-s1">argv</span>  <span class="pl-c"># Unused</span>
  <span class="pl-k">with</span> <span class="pl-s1">maps</span>.<span class="pl-v">Mesh</span>(<span class="pl-s1">np</span>.<span class="pl-en">asarray</span>(<span class="pl-s1">jax</span>.<span class="pl-en">devices</span>()), <span class="pl-s1">axis_names</span><span class="pl-c1">=</span>[<span class="pl-s">'x'</span>]):
    <span class="pl-en">print</span>(<span class="pl-en">f</span>().<span class="pl-s1">shape</span>)
    <span class="pl-en">print</span>(<span class="pl-s1">jax</span>.<span class="pl-en">eval_shape</span>(<span class="pl-s1">f</span>))

<span class="pl-k">if</span> <span class="pl-s1">__name__</span> <span class="pl-c1">==</span> <span class="pl-s">'__main__'</span>:
  <span class="pl-s1">app</span>.<span class="pl-en">run</span>(<span class="pl-s1">main</span>)</pre></div>
<h2 dir="auto">Expected output:</h2>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="(32, 10)
ShapeDtypeStruct(shape=(32, 10), dtype=float32)"><pre class="notranslate"><code class="notranslate">(32, 10)
ShapeDtypeStruct(shape=(32, 10), dtype=float32)
</code></pre></div>
<h2 dir="auto">Observed output:</h2>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="(32, 10)
ShapeDtypeStruct(shape=(16, 10), dtype=float32)"><pre class="notranslate"><code class="notranslate">(32, 10)
ShapeDtypeStruct(shape=(16, 10), dtype=float32)
</code></pre></div>
<p dir="auto">More generally, it would be nice if <code class="notranslate">eval_shape()</code> could recognise <code class="notranslate">GlobalDeviceArray</code> and produce a subtype of <code class="notranslate">ShapeDtypeStruct</code> conveying not just global shape but also the shard index and replica information derived from the corresponding <code class="notranslate">PartitionSpec</code>. Otherwise it is currently very hard to work out what shape a <code class="notranslate">GlobalDeviceArray</code> is going to have and what shape its shards will have.</p>