Merge pull request #399 from hawkinsp/types

Add {float16,uint16,uint8,int16,int8} types to abstract_arrays.
google · Feb 18, 2019 · 99abdf9 · 99abdf9
2 parents 848b769 + 0129e94
commit 99abdf9
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/jax/abstract_arrays.py b/jax/abstract_arrays.py
@@ -161,9 +161,11 @@ def zeros_like_array(x):
   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