Update ag4masses/alphageometry/decoder_stack.py

ag4masses/alphageometry/decoder_stack.py

@@ -1,55 +1,55 @@
-# Copyright 2023 DeepMind Technologies Limited
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#    http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-
-"""The decoder stack in inference mode."""
-
-from typing import Any, Tuple
-
-import gin
-from meliad_lib.meliad.transformer import decoder_stack
-import transformer_layer as tl
-
-
-struct = decoder_stack.struct
-nn_components = decoder_stack.nn_components
-position = decoder_stack.position
-jnp = decoder_stack.jnp
-attention = decoder_stack.attention
-
-DStackWindowState = decoder_stack.DStackWindowState
-
-Array = Any
-
-TransformerTaskConfig = decoder_stack.TransformerTaskConfig
-
-DStackDecoderState = Tuple[tl.DecoderState, ...]
-
-
-@gin.configurable
-class DecoderStackGenerate(decoder_stack.DecoderStack):
-  """Stack of transformer decoder layers."""
-
-  layer_factory = tl.TransformerLayerGenerate
-
-  def init_decoder_state_vanilla(
-      self, sequence_length: int, start_of_sequence: Array
-  ) -> DStackDecoderState:
-    """Return initial state for autoregressive generation."""
-    return tuple(
-        [
-            layer.init_decoder_state_vanilla(sequence_length, start_of_sequence)
-            for layer in self.transformer_layers
-        ]
-    )
+# Copyright 2023 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""The decoder stack in inference mode."""
+
+from typing import Any, Tuple
+
+import gin
+from aglib.meliad.transformer import decoder_stack
+import transformer_layer as tl
+
+
+struct = decoder_stack.struct
+nn_components = decoder_stack.nn_components
+position = decoder_stack.position
+jnp = decoder_stack.jnp
+attention = decoder_stack.attention
+
+DStackWindowState = decoder_stack.DStackWindowState
+
+Array = Any
+
+TransformerTaskConfig = decoder_stack.TransformerTaskConfig
+
+DStackDecoderState = Tuple[tl.DecoderState, ...]
+
+
+@gin.configurable
+class DecoderStackGenerate(decoder_stack.DecoderStack):
+  """Stack of transformer decoder layers."""
+
+  layer_factory = tl.TransformerLayerGenerate
+
+  def init_decoder_state_vanilla(
+      self, sequence_length: int, start_of_sequence: Array
+  ) -> DStackDecoderState:
+    """Return initial state for autoregressive generation."""
+    return tuple(
+        [
+            layer.init_decoder_state_vanilla(sequence_length, start_of_sequence)
+            for layer in self.transformer_layers
+        ]
+    )