add with_sharding_constraint to attention

yangdsh · Mar 12, 2024 · 1a0b9a3 · 1a0b9a3
1 parent 47b21f0
commit 1a0b9a3
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/axlearn/common/attention.py b/axlearn/common/attention.py
@@ -94,6 +94,7 @@
     get_or_none,
     shapes,
     split_prng_key,
+    with_sharding_constraint
 )
 
 NEG_INF = -1e15