Change attention to pytype_library.

lingvo-bot · copybara-github · commit 36a1e3148645 · 2023-11-28T14:59:15.000-08:00
PiperOrigin-RevId: 586113221
diff --git a/lingvo/core/BUILD b/lingvo/core/BUILD
@@ -40,7 +40,7 @@ lingvo_cuda_py_test(
     ],
 )
 
-py_library(
+pytype_library(
     name = "attention",
     srcs = ["attention.py"],
     deps = [
diff --git a/lingvo/core/attention.py b/lingvo/core/attention.py
@@ -2713,7 +2713,7 @@ def ComputeContextVectorWithSource(self,
                                      theta,
                                      packed_src,
                                      query_vec,
-                                     attention_state,
+                                     attention_state=None,
                                      per_step_source_padding=None,
                                      query_segment_id=None):
     """Computes the context vector given the current query output.
@@ -2971,7 +2971,7 @@ def ComputeContextVectorWithSource(self,
                                      theta,
                                      packed_src,
                                      query_vec,
-                                     attention_state,
+                                     attention_state=None,
                                      per_step_source_padding=None,
                                      query_segment_id=None):
     """Computes the context vector given the current query output.
