Re-enable sharded_test.py, underlying JAX issue fixed.

END_PUBLIC *** Original change description *** BEGIN_PUBLIC Disable sharded_test.py in JAX backend. PiperOrigin-RevId: 395489287
ENate · Sep 8, 2021 · ca74f14 · ca74f14
1 parent e8c5712
commit ca74f14
Showing 1 changed file with 0 additions and 3 deletions.
diff --git a/tensorflow_probability/python/experimental/distribute/sharded_test.py b/tensorflow_probability/python/experimental/distribute/sharded_test.py
@@ -108,9 +108,6 @@ def run(key):
           continue
         self.assertNotEqual(sample[i], sample[j])
 
-  @test_util.disable_test_for_backend(
-      disable_jax=True,
-      reason='Throws NotImplementedError in JAX backend b/199209907')
   def test_kahan_custom_grad(self):
 
     def model_fn():