title: jax woes (on an NVDIA DGX box, no less)
body: <p>I am trying to run jax on an nvidia dgx box, but am failing miserably, thus:</p>
  <pre><code>&gt;&gt;&gt; import jax
  &gt;&gt;&gt; import jax.numpy as jnp
  &gt;&gt;&gt; x = jnp.arange(10)
  2021-10-25 13:00:05.863667: W 
  external/org_tensorflow/tensorflow/stream_executor/gpu/asm_compiler.cc:80] Couldn't 
  get ptxas version string: INTERNAL: Couldn't invoke ptxas --version
  2021-10-25 13:00:05.864713: F 
  external/org_tensorflow/tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc:435] 
  ptxas returned an error during compilation of ptx to sass: 'INTERNAL: Failed to 
  launch ptxas'  If the error message indicates that a file could not be written, 
  please verify that sufficient filesystem space is provided.
  Aborted (core dumped)
  </code></pre>
  <p>Any suggestions would be much appreciated.</p>