Replace usage of tff.jax_computation with tff.jax.computation.

PiperOrigin-RevId: 625737178
google-parfait · Apr 17, 2024 · db481ad · db481ad
1 parent 3660e7e
commit db481ad
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tensorflow_federated/python/core/environments/jax_frontend/jax_computation.py b/tensorflow_federated/python/core/environments/jax_frontend/jax_computation.py
@@ -113,7 +113,7 @@ def _jax_wrapper_fn(
   Example:
 
   ```
-  @tff.jax_computation(np.int32)
+  @tff.jax.computation(np.int32)
   def comp(x):
     return jax.numpy.add(x, np.int32(10))
   ```