<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 issues.</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:</li>
</ul>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import jax.tools.colab_tpu
jax.tools.colab_tpu.setup_tpu()
jax.local_devices()"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">jax</span>.<span class="pl-s1">tools</span>.<span class="pl-s1">colab_tpu</span>
<span class="pl-s1">jax</span>.<span class="pl-s1">tools</span>.<span class="pl-s1">colab_tpu</span>.<span class="pl-en">setup_tpu</span>()
<span class="pl-s1">jax</span>.<span class="pl-en">local_devices</span>()</pre></div>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> If applicable, include full error messages/tracebacks.</li>
</ul>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="RuntimeError                              Traceback (most recent call last)
&lt;ipython-input-2-abdb42bb8a9a&gt; in &lt;module&gt;()
      3 import jax.tools.colab_tpu
      4 jax.tools.colab_tpu.setup_tpu()
----&gt; 5 jax.local_devices()

2 frames
/usr/local/lib/python3.7/dist-packages/jax/_src/lib/xla_bridge.py in _get_backend_uncached(platform)
    246     if backend is None:
    247       if platform in _backends_errors:
--&gt; 248         raise RuntimeError(f&quot;Requested backend {platform}, but it failed &quot;
    249                            f&quot;to initialize: {_backends_errors[platform]}&quot;)
    250       raise RuntimeError(f&quot;Unknown backend {platform}&quot;)

RuntimeError: Requested backend tpu_driver, but it failed to initialize: DEADLINE_EXCEEDED: Failed to connect to remote server at address: grpc://10.91.25.74:8470. Error from gRPC: Deadline Exceeded. Details: "><pre class="notranslate"><code class="notranslate">RuntimeError                              Traceback (most recent call last)
&lt;ipython-input-2-abdb42bb8a9a&gt; in &lt;module&gt;()
      3 import jax.tools.colab_tpu
      4 jax.tools.colab_tpu.setup_tpu()
----&gt; 5 jax.local_devices()

2 frames
/usr/local/lib/python3.7/dist-packages/jax/_src/lib/xla_bridge.py in _get_backend_uncached(platform)
    246     if backend is None:
    247       if platform in _backends_errors:
--&gt; 248         raise RuntimeError(f"Requested backend {platform}, but it failed "
    249                            f"to initialize: {_backends_errors[platform]}")
    250       raise RuntimeError(f"Unknown backend {platform}")

RuntimeError: Requested backend tpu_driver, but it failed to initialize: DEADLINE_EXCEEDED: Failed to connect to remote server at address: grpc://10.91.25.74:8470. Error from gRPC: Deadline Exceeded. Details: 
</code></pre></div>
<p dir="auto">On the other hand, using an earlier version of driver in <code class="notranslate">setup_tpu</code> <em>does</em> work, e.g.,</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="def setup_tpu():
  &quot;&quot;&quot;Sets up Colab to run on TPU.
  Note: make sure the Colab Runtime is set to Accelerator: TPU.
  &quot;&quot;&quot;
  global TPU_DRIVER_MODE

  if not TPU_DRIVER_MODE:
    colab_tpu_addr = os.environ['COLAB_TPU_ADDR'].split(':')[0]
    url = f'http://{colab_tpu_addr}:8475/requestversion/tpu_driver0.1-dev20211030'
    requests.post(url)
    TPU_DRIVER_MODE = 1

  # The following is required to use TPU Driver as JAX's backend.
  config.FLAGS.jax_xla_backend = &quot;tpu_driver&quot;
  config.FLAGS.jax_backend_target = &quot;grpc://&quot; + os.environ['COLAB_TPU_ADDR']"><pre class="notranslate"><code class="notranslate">def setup_tpu():
  """Sets up Colab to run on TPU.
  Note: make sure the Colab Runtime is set to Accelerator: TPU.
  """
  global TPU_DRIVER_MODE

  if not TPU_DRIVER_MODE:
    colab_tpu_addr = os.environ['COLAB_TPU_ADDR'].split(':')[0]
    url = f'http://{colab_tpu_addr}:8475/requestversion/tpu_driver0.1-dev20211030'
    requests.post(url)
    TPU_DRIVER_MODE = 1

  # The following is required to use TPU Driver as JAX's backend.
  config.FLAGS.jax_xla_backend = "tpu_driver"
  config.FLAGS.jax_backend_target = "grpc://" + os.environ['COLAB_TPU_ADDR']
</code></pre></div>
<p dir="auto">Hello,</p>
<p dir="auto">After reviewing the following code<br>
<a href="https://github.com/google/jax/blob/b4ec72deaefa1094891829cd0b78efb5cc19c28b/jax/_src/tree_util.py#LL545-L571C4">https://github.com/google/jax/blob/b4ec72deaefa1094891829cd0b78efb5cc19c28b/jax/_src/tree_util.py#LL545-L571C4</a><br>
If I understood correctly, if the user defines <code class="notranslate">flatten_func_with_keys</code> , <code class="notranslate">register_pytree_with_keys</code> will generate <code class="notranslate">flatten_func</code> registered with the non-key <code class="notranslate">jax</code> registry.</p>
<p dir="auto">I think, In some scenarios, the overhead from the list comprehension required for conversion from <code class="notranslate">flatten_func_with_keys</code> to <code class="notranslate">flatten_func</code> can not be justified by the usability of registering the tree once.</p>
<p dir="auto">Is it possible to decouple registering the <code class="notranslate">flatten_func_with_keys</code> from <code class="notranslate">flatten_func</code> inside <code class="notranslate">register_pytree_with_keys</code></p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="  def flatten_func(tree):
    key_children, treedef = flatten_with_keys(tree)
    return [c for _, c in key_children], treedef"><pre class="notranslate">  <span class="pl-k">def</span> <span class="pl-en">flatten_func</span>(<span class="pl-s1">tree</span>):
    <span class="pl-s1">key_children</span>, <span class="pl-s1">treedef</span> <span class="pl-c1">=</span> <span class="pl-en">flatten_with_keys</span>(<span class="pl-s1">tree</span>)
    <span class="pl-k">return</span> [<span class="pl-s1">c</span> <span class="pl-k">for</span> <span class="pl-s1">_</span>, <span class="pl-s1">c</span> <span class="pl-c1">in</span> <span class="pl-s1">key_children</span>], <span class="pl-s1">treedef</span></pre></div>