Remove type restrictions (google#4009)

dynamic_slice and dynamic_update_slice work with S16 types on TPU.
majnemer · Aug 10, 2020 · eb086f9 · eb086f9
1 parent d551cec
commit eb086f9
Showing 1 changed file with 0 additions and 2 deletions.
diff --git a/tests/lax_test.py b/tests/lax_test.py
@@ -1753,15 +1753,13 @@ def testReshapeWithUnusualShapes(self):
       "Shapes must be 1D sequences of concrete values of integer type.*",
       lambda: lax.reshape(np.ones(3,), (1.5, 2.0)))
 
-  @jtu.skip_on_devices("tpu")  # S16 not supported on TPU
   def testDynamicSliceTypeErrors(self):
     self.assertRaisesRegex(
       TypeError,
       "index arguments to dynamic_slice must be integers of the same type",
       lambda: lax.dynamic_slice(np.ones((3, 4), dtype=np.float32),
                                 (np.int32(1), np.int16(2)), (2, 2)))
 
-  @jtu.skip_on_devices("tpu")  # S16 not supported on TPU
   def testDynamicUpdateSliceTypeErrors(self):
     self.assertRaisesRegex(
       TypeError,