We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
float16 support
Add support for `np.float16`.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `jax/abstract_arrays.py`
Content:
```
1 # Copyright 2018 Google LLC
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     https://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 from __future__ import absolute_import
16 from __future__ import division
17 from __future__ import print_function
18 
19 import numpy as onp
20 import six
21 
22 from . import core
23 from . import ad_util
24 from . util import prod
25 from .lib import xla_bridge
26 
27 
28 def concretization_err_msg(fun):
29   fname = getattr(fun, "__name__", fun)
30   msg = ("Abstract value passed to `{}`, which requires a concrete value. "
31          "The function to be transformed can't be traced at the required level "
32          "of abstraction. If using `jit`, try using `static_argnums` or "
33          "applying `jit` to smaller subfunctions instead.")
34   return msg.format(fname)
35 
36 def concretization_function_error(fun):
37   def error(self, *args):
38     raise TypeError(concretization_err_msg(fun))
39   return error
40 
41 
42 class UnshapedArray(core.AbstractValue):
43   __slots__ = ['dtype']
44   array_abstraction_level = 3
45 
46   def __init__(self, dtype):
47     self.dtype = dtype
48 
49   def __eq__(self, other):
50     return type(self) is type(other) and self.dtype == other.dtype
51 
52   def __hash__(self):
53     return hash(str(self.dtype))
54 
55   def __repr__(self):
56     return '{}({})'.format(self.__class__.__name__, self.str_short())
57 
58   _bool = _nonzero = concretization_function_error(bool)
59   _float   = concretization_function_error(float)
60   _int     = concretization_function_error(int)
61   if six.PY2:
62     _long    = concretization_function_error(long)  # noqa: F821
63   _complex = concretization_function_error(complex)
64   _hex     = concretization_function_error(hex)
65   _oct     = concretization_function_error(oct)
66 
67   def at_least_vspace(self):
68     return self
69 
70   def join(self, other):
71     return self
72 
73   def str_short(self):
74     return onp.dtype(self.dtype).name
75 
76 
77 class ShapedArray(UnshapedArray):
78   __slots__ = ['shape']
79   array_abstraction_level = 2
80 
81   def __init__(self, shape, dtype):
82     self.dtype = onp.dtype(xla_bridge.canonicalize_dtype(dtype))
83     self.shape = shape
84 
85   ndim = property(lambda self: len(self.shape))
86   size = property(lambda self: prod(self.shape))
87 
88   def __eq__(self, other):
89     return (type(self) is type(other)
90             and self.dtype == other.dtype and self.shape == other.shape)
91 
92   def __hash__(self):
93     return hash((self.shape, str(self.dtype)))
94 
95   def at_least_vspace(self):
96     return self
97 
98   def join(self, other):
99     if self.shape == other.shape and self.dtype == other.dtype:
100       return self
101     elif self.dtype == other.dtype:
102       return UnshapedArray(self.dtype)
103     else:
104       raise TypeError(other)
105 
106   def str_short(self):
107     dtypestr = onp.dtype(self.dtype).name
108     shapestr = ','.join(map(str, self.shape))
109     return '{}[{}]'.format(dtypestr, shapestr)
110 
111   def __len__(self):
112     try:
113       return self.shape[0]
114     except IndexError:
115       raise TypeError("len() of unsized object")  # same as numpy error
116 
117   def _len(self, ignored_tracer):
118     return len(self)
119 
120 
121 class ConcreteArray(ShapedArray):
122   __slots__ = ['val']
123   array_abstraction_level = 0
124 
125   def __init__(self, val):
126     self.val = val
127     self.shape = onp.shape(val)
128     # canonicalized self.dtype doesn't necessarily match self.val
129     self.dtype = onp.dtype(xla_bridge.canonicalize_dtype(onp.result_type(val)))
130     assert self.dtype != onp.dtype('O')
131 
132   def __eq__(self, other):
133     return (type(self) is type(other) and self.dtype == other.dtype
134             and self.shape == other.shape and onp.all(self.val == other.val))
135 
136   def __hash__(self):
137     return id(self.val)
138 
139   def at_least_vspace(self):
140     return ShapedArray(self.shape, self.dtype)
141 
142   def join(self, other):
143     if self == other:
144       return self
145     elif self.shape == other.shape and self.dtype == other.dtype:
146       return ShapedArray(self.shape, self.dtype)
147     elif self.dtype == other.dtype:
148       return UnshapedArray(self.dtype)
149     else:
150       raise TypeError(other)
151 
152   def str_short(self):
153     return str(self.val)
154 
155 
156 def make_shaped_array(x):
157   dtype = xla_bridge.canonicalize_dtype(onp.result_type(x))
158   return ShapedArray(onp.shape(x), dtype)
159 
160 def zeros_like_array(x):
161   dtype = xla_bridge.canonicalize_dtype(onp.result_type(x))
162   return onp.broadcast_to(onp.array(0, dtype), onp.shape(x))
163 
164 array_types = [onp.ndarray, onp.float64, onp.float32, onp.complex64,
165                onp.complex128, onp.int64, onp.int32, onp.bool_, onp.uint64,
166                onp.uint32, complex, float, int, bool]
167 
168 for t in array_types:
169   core.pytype_aval_mappings[t] = ConcreteArray
170   ad_util.jaxval_zeros_likers[t] = zeros_like_array
171
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/jax/abstract_arrays.py b/jax/abstract_arrays.py
--- a/jax/abstract_arrays.py
+++ b/jax/abstract_arrays.py
@@ -161,9 +161,11 @@
   dtype = xla_bridge.canonicalize_dtype(onp.result_type(x))
   return onp.broadcast_to(onp.array(0, dtype), onp.shape(x))
 
-array_types = [onp.ndarray, onp.float64, onp.float32, onp.complex64,
-               onp.complex128, onp.int64, onp.int32, onp.bool_, onp.uint64,
-               onp.uint32, complex, float, int, bool]
+array_types = [onp.ndarray, onp.float64, onp.float32, onp.float16,
+               onp.complex64, onp.complex128,
+               onp.int64, onp.int32, onp.int16, onp.int8,
+               onp.bool_, onp.uint64, onp.uint32, onp.uint16, onp.uint8,
+               complex, float, int, bool]
 
 for t in array_types:
   core.pytype_aval_mappings[t] = ConcreteArray

{"golden_diff": "diff --git a/jax/abstract_arrays.py b/jax/abstract_arrays.py\n--- a/jax/abstract_arrays.py\n+++ b/jax/abstract_arrays.py\n@@ -161,9 +161,11 @@\n   dtype = xla_bridge.canonicalize_dtype(onp.result_type(x))\n   return onp.broadcast_to(onp.array(0, dtype), onp.shape(x))\n \n-array_types = [onp.ndarray, onp.float64, onp.float32, onp.complex64,\n-               onp.complex128, onp.int64, onp.int32, onp.bool_, onp.uint64,\n-               onp.uint32, complex, float, int, bool]\n+array_types = [onp.ndarray, onp.float64, onp.float32, onp.float16,\n+               onp.complex64, onp.complex128,\n+               onp.int64, onp.int32, onp.int16, onp.int8,\n+               onp.bool_, onp.uint64, onp.uint32, onp.uint16, onp.uint8,\n+               complex, float, int, bool]\n \n for t in array_types:\n   core.pytype_aval_mappings[t] = ConcreteArray\n", "issue": "float16 support\nAdd support for `np.float16`.\n", "before_files": [{"content": "# Copyright 2018 Google LLC\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport numpy as onp\nimport six\n\nfrom . import core\nfrom . import ad_util\nfrom . util import prod\nfrom .lib import xla_bridge\n\n\ndef concretization_err_msg(fun):\n  fname = getattr(fun, \"__name__\", fun)\n  msg = (\"Abstract value passed to `{}`, which requires a concrete value. \"\n         \"The function to be transformed can't be traced at the required level \"\n         \"of abstraction. If using `jit`, try using `static_argnums` or \"\n         \"applying `jit` to smaller subfunctions instead.\")\n  return msg.format(fname)\n\ndef concretization_function_error(fun):\n  def error(self, *args):\n    raise TypeError(concretization_err_msg(fun))\n  return error\n\n\nclass UnshapedArray(core.AbstractValue):\n  __slots__ = ['dtype']\n  array_abstraction_level = 3\n\n  def __init__(self, dtype):\n    self.dtype = dtype\n\n  def __eq__(self, other):\n    return type(self) is type(other) and self.dtype == other.dtype\n\n  def __hash__(self):\n    return hash(str(self.dtype))\n\n  def __repr__(self):\n    return '{}({})'.format(self.__class__.__name__, self.str_short())\n\n  _bool = _nonzero = concretization_function_error(bool)\n  _float   = concretization_function_error(float)\n  _int     = concretization_function_error(int)\n  if six.PY2:\n    _long    = concretization_function_error(long)  # noqa: F821\n  _complex = concretization_function_error(complex)\n  _hex     = concretization_function_error(hex)\n  _oct     = concretization_function_error(oct)\n\n  def at_least_vspace(self):\n    return self\n\n  def join(self, other):\n    return self\n\n  def str_short(self):\n    return onp.dtype(self.dtype).name\n\n\nclass ShapedArray(UnshapedArray):\n  __slots__ = ['shape']\n  array_abstraction_level = 2\n\n  def __init__(self, shape, dtype):\n    self.dtype = onp.dtype(xla_bridge.canonicalize_dtype(dtype))\n    self.shape = shape\n\n  ndim = property(lambda self: len(self.shape))\n  size = property(lambda self: prod(self.shape))\n\n  def __eq__(self, other):\n    return (type(self) is type(other)\n            and self.dtype == other.dtype and self.shape == other.shape)\n\n  def __hash__(self):\n    return hash((self.shape, str(self.dtype)))\n\n  def at_least_vspace(self):\n    return self\n\n  def join(self, other):\n    if self.shape == other.shape and self.dtype == other.dtype:\n      return self\n    elif self.dtype == other.dtype:\n      return UnshapedArray(self.dtype)\n    else:\n      raise TypeError(other)\n\n  def str_short(self):\n    dtypestr = onp.dtype(self.dtype).name\n    shapestr = ','.join(map(str, self.shape))\n    return '{}[{}]'.format(dtypestr, shapestr)\n\n  def __len__(self):\n    try:\n      return self.shape[0]\n    except IndexError:\n      raise TypeError(\"len() of unsized object\")  # same as numpy error\n\n  def _len(self, ignored_tracer):\n    return len(self)\n\n\nclass ConcreteArray(ShapedArray):\n  __slots__ = ['val']\n  array_abstraction_level = 0\n\n  def __init__(self, val):\n    self.val = val\n    self.shape = onp.shape(val)\n    # canonicalized self.dtype doesn't necessarily match self.val\n    self.dtype = onp.dtype(xla_bridge.canonicalize_dtype(onp.result_type(val)))\n    assert self.dtype != onp.dtype('O')\n\n  def __eq__(self, other):\n    return (type(self) is type(other) and self.dtype == other.dtype\n            and self.shape == other.shape and onp.all(self.val == other.val))\n\n  def __hash__(self):\n    return id(self.val)\n\n  def at_least_vspace(self):\n    return ShapedArray(self.shape, self.dtype)\n\n  def join(self, other):\n    if self == other:\n      return self\n    elif self.shape == other.shape and self.dtype == other.dtype:\n      return ShapedArray(self.shape, self.dtype)\n    elif self.dtype == other.dtype:\n      return UnshapedArray(self.dtype)\n    else:\n      raise TypeError(other)\n\n  def str_short(self):\n    return str(self.val)\n\n\ndef make_shaped_array(x):\n  dtype = xla_bridge.canonicalize_dtype(onp.result_type(x))\n  return ShapedArray(onp.shape(x), dtype)\n\ndef zeros_like_array(x):\n  dtype = xla_bridge.canonicalize_dtype(onp.result_type(x))\n  return onp.broadcast_to(onp.array(0, dtype), onp.shape(x))\n\narray_types = [onp.ndarray, onp.float64, onp.float32, onp.complex64,\n               onp.complex128, onp.int64, onp.int32, onp.bool_, onp.uint64,\n               onp.uint32, complex, float, int, bool]\n\nfor t in array_types:\n  core.pytype_aval_mappings[t] = ConcreteArray\n  ad_util.jaxval_zeros_likers[t] = zeros_like_array\n", "path": "jax/abstract_arrays.py"}], "after_files": [{"content": "# Copyright 2018 Google LLC\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport numpy as onp\nimport six\n\nfrom . import core\nfrom . import ad_util\nfrom . util import prod\nfrom .lib import xla_bridge\n\n\ndef concretization_err_msg(fun):\n  fname = getattr(fun, \"__name__\", fun)\n  msg = (\"Abstract value passed to `{}`, which requires a concrete value. \"\n         \"The function to be transformed can't be traced at the required level \"\n         \"of abstraction. If using `jit`, try using `static_argnums` or \"\n         \"applying `jit` to smaller subfunctions instead.\")\n  return msg.format(fname)\n\ndef concretization_function_error(fun):\n  def error(self, *args):\n    raise TypeError(concretization_err_msg(fun))\n  return error\n\n\nclass UnshapedArray(core.AbstractValue):\n  __slots__ = ['dtype']\n  array_abstraction_level = 3\n\n  def __init__(self, dtype):\n    self.dtype = dtype\n\n  def __eq__(self, other):\n    return type(self) is type(other) and self.dtype == other.dtype\n\n  def __hash__(self):\n    return hash(str(self.dtype))\n\n  def __repr__(self):\n    return '{}({})'.format(self.__class__.__name__, self.str_short())\n\n  _bool = _nonzero = concretization_function_error(bool)\n  _float   = concretization_function_error(float)\n  _int     = concretization_function_error(int)\n  if six.PY2:\n    _long    = concretization_function_error(long)  # noqa: F821\n  _complex = concretization_function_error(complex)\n  _hex     = concretization_function_error(hex)\n  _oct     = concretization_function_error(oct)\n\n  def at_least_vspace(self):\n    return self\n\n  def join(self, other):\n    return self\n\n  def str_short(self):\n    return onp.dtype(self.dtype).name\n\n\nclass ShapedArray(UnshapedArray):\n  __slots__ = ['shape']\n  array_abstraction_level = 2\n\n  def __init__(self, shape, dtype):\n    self.dtype = onp.dtype(xla_bridge.canonicalize_dtype(dtype))\n    self.shape = shape\n\n  ndim = property(lambda self: len(self.shape))\n  size = property(lambda self: prod(self.shape))\n\n  def __eq__(self, other):\n    return (type(self) is type(other)\n            and self.dtype == other.dtype and self.shape == other.shape)\n\n  def __hash__(self):\n    return hash((self.shape, str(self.dtype)))\n\n  def at_least_vspace(self):\n    return self\n\n  def join(self, other):\n    if self.shape == other.shape and self.dtype == other.dtype:\n      return self\n    elif self.dtype == other.dtype:\n      return UnshapedArray(self.dtype)\n    else:\n      raise TypeError(other)\n\n  def str_short(self):\n    dtypestr = onp.dtype(self.dtype).name\n    shapestr = ','.join(map(str, self.shape))\n    return '{}[{}]'.format(dtypestr, shapestr)\n\n  def __len__(self):\n    try:\n      return self.shape[0]\n    except IndexError:\n      raise TypeError(\"len() of unsized object\")  # same as numpy error\n\n  def _len(self, ignored_tracer):\n    return len(self)\n\n\nclass ConcreteArray(ShapedArray):\n  __slots__ = ['val']\n  array_abstraction_level = 0\n\n  def __init__(self, val):\n    self.val = val\n    self.shape = onp.shape(val)\n    # canonicalized self.dtype doesn't necessarily match self.val\n    self.dtype = onp.dtype(xla_bridge.canonicalize_dtype(onp.result_type(val)))\n    assert self.dtype != onp.dtype('O')\n\n  def __eq__(self, other):\n    return (type(self) is type(other) and self.dtype == other.dtype\n            and self.shape == other.shape and onp.all(self.val == other.val))\n\n  def __hash__(self):\n    return id(self.val)\n\n  def at_least_vspace(self):\n    return ShapedArray(self.shape, self.dtype)\n\n  def join(self, other):\n    if self == other:\n      return self\n    elif self.shape == other.shape and self.dtype == other.dtype:\n      return ShapedArray(self.shape, self.dtype)\n    elif self.dtype == other.dtype:\n      return UnshapedArray(self.dtype)\n    else:\n      raise TypeError(other)\n\n  def str_short(self):\n    return str(self.val)\n\n\ndef make_shaped_array(x):\n  dtype = xla_bridge.canonicalize_dtype(onp.result_type(x))\n  return ShapedArray(onp.shape(x), dtype)\n\ndef zeros_like_array(x):\n  dtype = xla_bridge.canonicalize_dtype(onp.result_type(x))\n  return onp.broadcast_to(onp.array(0, dtype), onp.shape(x))\n\narray_types = [onp.ndarray, onp.float64, onp.float32, onp.float16,\n               onp.complex64, onp.complex128,\n               onp.int64, onp.int32, onp.int16, onp.int8,\n               onp.bool_, onp.uint64, onp.uint32, onp.uint16, onp.uint8,\n               complex, float, int, bool]\n\nfor t in array_types:\n  core.pytype_aval_mappings[t] = ConcreteArray\n  ad_util.jaxval_zeros_likers[t] = zeros_like_array\n", "path": "jax/abstract_arrays.py"}]}