Merge pull request #405 from hawkinsp/fixes

Remove obsolete workarounds for bugs that seem fixed.

jax/interpreters/xla.py

@@ -262,17 +262,7 @@ def __init__(self, device_buffer, shape, dtype, ndim, size):
   def _value(self):
     if self._npy_value is None:
       self._npy_value = self.device_buffer.to_py()
-      try:
-        self._npy_value.flags.writeable = False
-      except AttributeError:
-        # TODO(mattjj): bug with C64 on TPU backend, C64 values returned as pair
-        if onp.issubdtype(self.dtype, onp.complexfloating):
-          a, b = self._npy_value
-          npy_value = onp.stack([a, b], -1).view(self.dtype).reshape(self.shape)
-          npy_value.flags.writeable = False
-          self._npy_value = npy_value
-        else:
-          raise
+      self._npy_value.flags.writeable = False
     return self._npy_value
 
   def copy(self):

tests/lax_numpy_test.py

@@ -1243,10 +1243,6 @@ def testRoll(self, shape, dtype, shifts, axis, rng):
       for rng_indices in [jtu.rand_int(-5, 5)]))
   def testTake(self, shape, dtype, index_shape, index_dtype, axis, mode, rng,
                rng_indices):
-    if (FLAGS.jax_test_dut.startswith("tpu")
-        and onp.issubdtype(dtype, onp.complexfloating)):
-      self.skipTest("skipping complex dtype on TPU")  # TODO(mattjj): investigate failures
-
     def args_maker():
       x = rng(shape, dtype)
       i = rng_indices(index_shape, index_dtype)