Repository: google/jax

Issue:
JAX with GPU support requires jaxlib version 0.1.26
I've got an issue.

So, I followed the installation instructions here:
![image](https://user-images.githubusercontent.com/12001304/64115427-e8b98300-cd8f-11e9-879c-b4935195991c.png)

Then I run the following line in `JupyterLab`:
```python
import jax.numpy as jnp
```
And that produces the following error:
```python
ValueError: jaxlib is version 0.1.23, but this version of jax requires version 0.1.26.
```

I went back and changed the second last line to:
```bash
pip install --upgrade $BASE_URL/$CUDA_VERSION/jaxlib-0.1.26-$PYTHON_VERSION-none-$PLATFORM.whl
```
Tried running. But, was denied.
![image](https://user-images.githubusercontent.com/12001304/64115630-6bdad900-cd90-11e9-92d6-ebc6dc1c15a0.png)


Discussion:
Thanks! I think our documentation is out-of-date, as we forgot to update it when we bumped the minimum jaxlib version.

Resolving patch:
diff --git a/jax/lax_linalg.py b/jax/lax_linalg.py
--- a/jax/lax_linalg.py
+++ b/jax/lax_linalg.py
@@ -134,12 +134,7 @@ def _nan_like(c, operand):
     nan = c.Constant(onp.array(onp.nan, dtype=dtype))
   return c.Broadcast(nan, shape.dimensions())
 
-# TODO(phawkins): remove if-condition after increasing minimum Jaxlib version to
-# 0.1.23.
-if hasattr(lapack, "potrf"):
-  _cpu_potrf = lapack.potrf
-else:
-  _cpu_potrf = _unpack_tuple(lapack.jax_potrf, 2)
+_cpu_potrf = lapack.potrf
 
 def cholesky_cpu_translation_rule(c, operand):
   shape = c.GetShape(operand)
@@ -181,12 +176,7 @@ def eig_abstract_eval(operand):
     raise NotImplementedError
   return w, vl, vr
 
-# TODO(phawkins): remove if-condition after increasing minimum Jaxlib version to
-# 0.1.23.
-if hasattr(lapack, "geev"):
-  _cpu_geev = lapack.geev
-else:
-  _cpu_geev = _unpack_tuple(lapack.jax_geev, 4)
+_cpu_geev = lapack.geev
 
 def eig_cpu_translation_rule(c, operand):
   shape = c.GetShape(operand)
@@ -294,21 +284,13 @@ def eigh_batching_rule(batched_args, batch_dims, lower):
 xla.translations[eigh_p] = eigh_translation_rule
 ad.primitive_jvps[eigh_p] = eigh_jvp_rule
 
-# TODO(phawkins): remove if-condition after increasing minimum Jaxlib version to
-# 0.1.23.
-if hasattr(lapack, "syevd"):
-  _cpu_syevd = lapack.syevd
-else:
-  _cpu_syevd = _unpack_tuple(lapack.jax_syevd, 3)
+_cpu_syevd = lapack.syevd
 
 xla.backend_specific_translations['cpu'][eigh_p] = partial(
   _eigh_cpu_gpu_translation_rule, _cpu_syevd)
 
-# TODO(phawkins): remove if-condition after increasing minimum Jaxlib version to
-# 0.1.23.
-if cusolver:
-  xla.backend_specific_translations['gpu'][eigh_p] = partial(
-    _eigh_cpu_gpu_translation_rule, cusolver.syevd)
+xla.backend_specific_translations['gpu'][eigh_p] = partial(
+  _eigh_cpu_gpu_translation_rule, cusolver.syevd)
 batching.primitive_batchers[eigh_p] = eigh_batching_rule
 
 
@@ -612,12 +594,7 @@ def _lu_cpu_gpu_translation_rule(getrf_impl, c, operand):
 ad.primitive_jvps[lu_p] = _lu_jvp_rule
 batching.primitive_batchers[lu_p] = _lu_batching_rule
 
-# TODO(phawkins): remove if-condition after increasing minimum Jaxlib version to
-# 0.1.23.
-if hasattr(lapack, "getrf"):
-  _cpu_getrf = lapack.getrf
-else:
-  _cpu_getrf = _unpack_tuple(lapack.jax_getrf, 3)
+_cpu_getrf = lapack.getrf
 
 xla.backend_specific_translations['cpu'][lu_p] = partial(
   _lu_cpu_gpu_translation_rule, _cpu_getrf)
@@ -803,18 +780,10 @@ def svd_batching_rule(batched_args, batch_dims, full_matrices, compute_uv):
 batching.primitive_batchers[svd_p] = svd_batching_rule
 xla.translations[svd_p] = svd_translation_rule
 
-# TODO(phawkins): remove if-condition after increasing minimum Jaxlib version to
-# 0.1.23.
-if hasattr(lapack, "gesdd"):
-  _cpu_gesdd = lapack.gesdd
-else:
-  _cpu_gesdd = _unpack_tuple(lapack.jax_gesdd, 4)
+_cpu_gesdd = lapack.gesdd
 
 xla.backend_specific_translations['cpu'][svd_p] = partial(
   _svd_cpu_gpu_translation_rule, _cpu_gesdd)
 
-# TODO(phawkins): remove if-condition after increasing minimum Jaxlib version to
-# 0.1.23.
-if cusolver:
-  xla.backend_specific_translations['gpu'][svd_p] = partial(
-    _svd_cpu_gpu_translation_rule, cusolver.gesvd)
+xla.backend_specific_translations['gpu'][svd_p] = partial(
+  _svd_cpu_gpu_translation_rule, cusolver.gesvd)
diff --git a/jax/lib/__init__.py b/jax/lib/__init__.py
--- a/jax/lib/__init__.py
+++ b/jax/lib/__init__.py
@@ -41,16 +41,5 @@ def _check_jaxlib_version():
 from jaxlib import xrt
 from jaxlib import lapack
 
-# TODO(phawkins): make the import unconditional when the minimum Jaxlib version
-# has been increased to 0.1.23.
-try:
-  from jaxlib import pytree
-except ImportError:
-  pytree = None
-
-# TODO(phawkins): make the import unconditional when the minimum Jaxlib version
-# has been increased to 0.1.23.
-try:
-  from jaxlib import cusolver
-except ImportError:
-  cusolver = None
+from jaxlib import pytree
+from jaxlib import cusolver
diff --git a/jax/tree_util.py b/jax/tree_util.py
--- a/jax/tree_util.py
+++ b/jax/tree_util.py
@@ -46,308 +46,85 @@
 from .util import unzip2, partial, safe_map
 
 
-# TODO(phawkins): use the first case unconditionally when the minimum Jaxlib
-# version has been increased to 0.1.23.
-if pytree:
-
-  def tree_map(f, tree):
-    """Map a function over a pytree to produce a new pytree.
-
-    Args:
-      f: function to be applied at each leaf.
-      tree: a pytree to be mapped over.
-
-    Returns:
-      A new pytree with the same structure as `tree` but with the value at each
-      leaf given by `f(x)` where `x` is the value at the corresponding leaf in
-      `tree`.
-    """
-    leaves, treedef = pytree.flatten(tree)
-    return treedef.unflatten(map(f, leaves))
-
-  def tree_multimap(f, tree, *rest):
-    """Map a multi-input function over pytree args to produce a new pytree.
-
-    Args:
-      f: function that takes `1 + len(rest)` arguments, to be applied at the
-        corresponding leaves of the pytrees.
-      tree: a pytree to be mapped over, with each leaf providing the first
-        positional argument to `f`.
-      *rest: a tuple of pytrees, each of which has the same structure as tree or
-        or has tree as a prefix.
-    Returns:
-      A new pytree with the same structure as `tree` but with the value at each
-      leaf given by `f(x, *xs)` where `x` is the value at the corresponding leaf
-      in `tree` and `xs` is the tuple of values at corresponding nodes in
-      `rest`.
-    """
-    leaves, treedef = pytree.flatten(tree)
-    all_leaves = [leaves] + [treedef.flatten_up_to(r) for r in rest]
-    return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
-
-  def tree_leaves(tree):
-    return pytree.flatten(tree)[0]
-
-  def process_pytree(process_node, tree):
-    leaves, treedef = pytree.flatten(tree)
-    return treedef.walk(process_node, None, leaves), treedef
-
-  tree_flatten = pytree.flatten
-
-  def build_tree(treedef, xs):
-    return treedef.from_iterable_tree(xs)
-
-  def treedef_is_leaf(treedef):
-    return treedef.num_nodes == 1
-
-  def tree_unflatten(treedef, xs):
-    return treedef.unflatten(xs)
-
-  def tree_transpose(outer_treedef, inner_treedef, pytree_to_transpose):
-    flat, treedef = tree_flatten(pytree_to_transpose)
-    expected_treedef = outer_treedef.compose(inner_treedef)
-    if treedef != expected_treedef:
-      raise TypeError("Mismatch\n{}\n != \n{}".format(treedef, expected_treedef))
-
-    inner_size = inner_treedef.num_leaves
-    outer_size = outer_treedef.num_leaves
-    flat = iter(flat)
-    lol = [[next(flat) for _ in range(inner_size)] for __ in range(outer_size)]
-    transposed_lol = zip(*lol)
-    subtrees = map(partial(tree_unflatten, outer_treedef), transposed_lol)
-    return tree_unflatten(inner_treedef, subtrees)
-
-  def tree_structure(tree):
-    _, treedef = pytree.flatten(tree)
-    return treedef
-
-  def treedef_tuple(trees):
-    return pytree.tuple(list(trees))
-
-  def treedef_children(treedef):
-      return treedef.children()
-
-  register_pytree_node = pytree.register_node
-
-else:
-  def tree_map(f, tree):
-    """Map a function over a pytree to produce a new pytree.
-
-    Args:
-      f: function to be applied at each leaf.
-      tree: a pytree to be mapped over.
-
-    Returns:
-      A new pytree with the same structure as `tree` but with the value at each
-      leaf given by `f(x)` where `x` is the value at the corresponding leaf in
-      `tree`.
-    """
-    node_type = _get_node_type(tree)
-    if node_type:
-      children, node_spec = node_type.to_iterable(tree)
-      new_children = [tree_map(f, child) for child in children]
-      return node_type.from_iterable(node_spec, new_children)
-    else:
-      return f(tree)
-
-  def tree_multimap(f, tree, *rest):
-    """Map a multi-input function over pytree args to produce a new pytree.
-
-    Args:
-      f: function that takes `1 + len(rest)` arguments, to be applied at the
-        corresponding leaves of the pytrees.
-      tree: a pytree to be mapped over, with each leaf providing the first
-        positional argument to `f`.
-      *rest: a tuple of pytrees, each with the same structure as `tree`.
-
-    Returns:
-      A new pytree with the same structure as `tree` but with the value at each
-      leaf given by `f(x, *xs)` where `x` is the value at the corresponding leaf
-      in `tree` and `xs` is the tuple of values at corresponding leaves in `rest`.
-    """
-    node_type = _get_node_type(tree)
-    if node_type:
-      children, aux_data = node_type.to_iterable(tree)
-      all_children = [children]
-      for other_tree in rest:
-        other_node_type = _get_node_type(other_tree)
-        if node_type != other_node_type:
-          raise TypeError('Mismatch: {} != {}'.format(other_node_type, node_type))
-        other_children, other_aux_data = node_type.to_iterable(other_tree)
-        if other_aux_data != aux_data:
-          raise TypeError('Mismatch: {} != {}'.format(other_aux_data, aux_data))
-        all_children.append(other_children)
-
-      new_children = [tree_multimap(f, *xs) for xs in zip(*all_children)]
-      return node_type.from_iterable(aux_data, new_children)
-    else:
-      return f(tree, *rest)
-
-  def _walk_pytree(f_node, f_leaf, tree):
-    node_type = _get_node_type(tree)
-    if node_type:
-      children, node_spec = node_type.to_iterable(tree)
-      proc_children, child_specs = unzip2([_walk_pytree(f_node, f_leaf, child)
-                                           for child in children])
-      tree_def = _PyTreeDef(node_type, node_spec, child_specs)
-      return f_node(proc_children), tree_def
-    else:
-      return f_leaf(tree), leaf
-
-  def process_pytree(process_node, tree):
-    return _walk_pytree(process_node, lambda x: x, tree)
-
-  def build_tree(treedef, xs):
-    if treedef is leaf:
-      return xs
-    else:
-      # We use 'iter' for clearer error messages
-      children = safe_map(build_tree, iter(treedef.children), iter(xs))
-      return treedef.node_type.from_iterable(treedef.node_data, children)
-
-  def tree_leaves(tree):
-    """Generator that iterates over all leaves of a pytree."""
-    node_type = _get_node_type(tree)
-    if node_type:
-      children, _ = node_type.to_iterable(tree)
-      for child in children:
-        # TODO(mattjj,phawkins): use 'yield from' when PY2 is dropped
-        for leaf in tree_leaves(child):
-          yield leaf
-    else:
-      yield tree
-
-  def tree_flatten(tree):
-    itr, treedef = _walk_pytree(it.chain.from_iterable, lambda x: (x,), tree)
-    return list(itr), treedef
-
-  def _tree_unflatten(xs, treedef):
-    if treedef is leaf:
-      return next(xs)
-    else:
-      children = tuple(map(partial(_tree_unflatten, xs), treedef.children))
-      return treedef.node_type.from_iterable(treedef.node_data, children)
-
-  def tree_unflatten(treedef, xs):
-    return _tree_unflatten(iter(xs), treedef)
-
-  def tree_transpose(outer_treedef, inner_treedef, pytree_to_transpose):
-    flat, treedef = tree_flatten(pytree_to_transpose)
-    expected_treedef = _nested_treedef(inner_treedef, outer_treedef)
-    if treedef != expected_treedef:
-      raise TypeError("Mismatch\n{}\n != \n{}".format(treedef, expected_treedef))
-
-    inner_size = _num_leaves(inner_treedef)
-    outer_size = _num_leaves(outer_treedef)
-    flat = iter(flat)
-    lol = [[next(flat) for _ in range(inner_size)] for __ in range(outer_size)]
-    transposed_lol = zip(*lol)
-    subtrees = map(partial(tree_unflatten, outer_treedef), transposed_lol)
-    return tree_unflatten(inner_treedef, subtrees)
-
-  def _num_leaves(treedef):
-    return 1 if treedef is leaf else sum(map(_num_leaves, treedef.children))
-
-  def _nested_treedef(inner, outer):
-    # just used in tree_transpose error checking
-    if outer is leaf:
-      return inner
-    else:
-      children = map(partial(_nested_treedef, inner), outer.children)
-      return _PyTreeDef(outer.node_type, outer.node_data, tuple(children))
-
-  def tree_structure(tree):
-    _, spec = process_pytree(lambda _: None, tree)
-    return spec
-
-
-  class _PyTreeDef(object):
-    __slots__ = ("node_type", "node_data", "children")
-
-    def __init__(self, node_type, node_data, children):
-      self.node_type = node_type
-      self.node_data = node_data
-      self.children = children
-
-    def __repr__(self):
-      if self.node_data is None:
-        data_repr = ""
-      else:
-        data_repr = "[{}]".format(self.node_data)
-
-      return "PyTree({}{}, [{}])".format(self.node_type.name, data_repr,
-                                       ','.join(map(repr, self.children)))
-
-    def __hash__(self):
-      return hash((self.node_type, self.node_data, tuple(self.children)))
-
-    def __eq__(self, other):
-      if other is leaf:
-        return False
-      else:
-        return (self.node_type == other.node_type and
-                self.node_data == other.node_data and
-                self.children == other.children)
-
-    def __ne__(self, other):
-      return not self == other
-
-
-  class _PyLeaf(object):
-    __slots__ = ()
-
-    def __repr__(self):
-      return '*'
-
-  leaf = _PyLeaf()
-
-  def treedef_is_leaf(treedef):
-    return treedef is leaf
-
-  def treedef_tuple(treedefs):
-    return _PyTreeDef(node_types[tuple], None, tuple(treedefs))
-
-  def treedef_children(treedef):
-    return treedef.children
-
-  def dict_to_iterable(xs):
-    keys = tuple(sorted(xs.keys()))
-    return tuple(map(xs.get, keys)), keys
-
-  class NodeType(object):
-    def __init__(self, name, to_iterable, from_iterable):
-      self.name = name
-      self.to_iterable = to_iterable
-      self.from_iterable = from_iterable
-
-    def __repr__(self):
-      return self.name
-
-  node_types = {}
-
-  def register_pytree_node(py_type, to_iterable, from_iterable):
-    assert py_type not in node_types
-    node_types[py_type] = NodeType(str(py_type), to_iterable, from_iterable)
-
-  register_pytree_node(tuple, lambda xs: (xs, None), lambda _, xs: tuple(xs))
-  register_pytree_node(list, lambda xs: (tuple(xs), None), lambda _, xs: list(xs))
-  register_pytree_node(dict, dict_to_iterable, lambda keys, xs: dict(zip(keys, xs)))
-
-  # To handle namedtuples, we can't just use the standard table of node_types
-  # because every namedtuple creates its own type and thus would require its own
-  # entry in the table. Instead we use a heuristic check on the type itself to
-  # decide whether it's a namedtuple type, and if so treat it as a pytree node.
-  def _get_node_type(maybe_tree):
-    t = type(maybe_tree)
-    return node_types.get(t) or _namedtuple_node(t)
+def tree_map(f, tree):
+  """Map a function over a pytree to produce a new pytree.
 
-  def _namedtuple_node(t):
-    if issubclass(t, tuple) and hasattr(t, '_fields'):
-      return NamedtupleNode
+  Args:
+    f: function to be applied at each leaf.
+    tree: a pytree to be mapped over.
+
+  Returns:
+    A new pytree with the same structure as `tree` but with the value at each
+    leaf given by `f(x)` where `x` is the value at the corresponding leaf in
+    `tree`.
+  """
+  leaves, treedef = pytree.flatten(tree)
+  return treedef.unflatten(map(f, leaves))
+
+def tree_multimap(f, tree, *rest):
+  """Map a multi-input function over pytree args to produce a new pytree.
+
+  Args:
+    f: function that takes `1 + len(rest)` arguments, to be applied at the
+      corresponding leaves of the pytrees.
+    tree: a pytree to be mapped over, with each leaf providing the first
+      positional argument to `f`.
+    *rest: a tuple of pytrees, each of which has the same structure as tree or
+      or has tree as a prefix.
+  Returns:
+    A new pytree with the same structure as `tree` but with the value at each
+    leaf given by `f(x, *xs)` where `x` is the value at the corresponding leaf
+    in `tree` and `xs` is the tuple of values at corresponding nodes in
+    `rest`.
+  """
+  leaves, treedef = pytree.flatten(tree)
+  all_leaves = [leaves] + [treedef.flatten_up_to(r) for r in rest]
+  return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
+
+def tree_leaves(tree):
+  return pytree.flatten(tree)[0]
+
+def process_pytree(process_node, tree):
+  leaves, treedef = pytree.flatten(tree)
+  return treedef.walk(process_node, None, leaves), treedef
+
+tree_flatten = pytree.flatten
+
+def build_tree(treedef, xs):
+  return treedef.from_iterable_tree(xs)
+
+def treedef_is_leaf(treedef):
+  return treedef.num_nodes == 1
+
+def tree_unflatten(treedef, xs):
+  return treedef.unflatten(xs)
+
+def tree_transpose(outer_treedef, inner_treedef, pytree_to_transpose):
+  flat, treedef = tree_flatten(pytree_to_transpose)
+  expected_treedef = outer_treedef.compose(inner_treedef)
+  if treedef != expected_treedef:
+    raise TypeError("Mismatch\n{}\n != \n{}".format(treedef, expected_treedef))
+
+  inner_size = inner_treedef.num_leaves
+  outer_size = outer_treedef.num_leaves
+  flat = iter(flat)
+  lol = [[next(flat) for _ in range(inner_size)] for __ in range(outer_size)]
+  transposed_lol = zip(*lol)
+  subtrees = map(partial(tree_unflatten, outer_treedef), transposed_lol)
+  return tree_unflatten(inner_treedef, subtrees)
+
+def tree_structure(tree):
+  _, treedef = pytree.flatten(tree)
+  return treedef
+
+def treedef_tuple(trees):
+  return pytree.tuple(list(trees))
+
+def treedef_children(treedef):
+    return treedef.children()
+
+register_pytree_node = pytree.register_node
 
-  NamedtupleNode = NodeType('namedtuple',
-                            lambda xs: (tuple(xs), type(xs)),
-                            lambda t, xs: t(*xs))
 
 
 def tree_reduce(f, tree):


Test patch: