Export jax.lax.sharding_constraint_p

PiperOrigin-RevId: 534566582
jax-ml · May 23, 2023 · 7f7f995 · 7f7f995
1 parent db87167
commit 7f7f995
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/jax/lax/__init__.py b/jax/lax/__init__.py
@@ -369,6 +369,7 @@
 from jax.lax import linalg as linalg
 
 from jax._src.pjit import with_sharding_constraint as with_sharding_constraint
+from jax._src.pjit import sharding_constraint_p as sharding_constraint_p
 from jax._src.dispatch import device_put_p as device_put_p
 
 from math import prod  # TODO(phawkins): remove this accidental export