Add support for using basic block context to GRANITE models.

gematria/granite/python/graph_builder_model_base.py

@@ -58,6 +58,9 @@ class GraphBuilderModelBase(
       'GraphBuilderModelBase.instruction_node_mask'
   )
 
+  # The name of the input tensor that receives the context node mask.
+  CONTEXT_NODE_MASK_TENSOR_NAME = 'GraphBuilderModelBase.context_node_mask'
+
   # The name of the input tensor that holds the instruction annotations.
   INSTRUCTION_ANNOTATIONS_TENSOR_NAME = (
       'GraphBuilderModelBase.instruction_annotations'
@@ -75,6 +78,12 @@ class GraphBuilderModelBase(
   # further processing during readout.
   _instruction_node_mask: tf.Tensor
 
+  # A Boolean tensor placeholder that receives a mask for context nodes of the
+  # same shape as `_instruction_node_mask`. A given element is True if the
+  # corresponding node is an instruction belonging to either the back or front
+  # context. The mask is used to exclude context nodes from predictions.
+  _context_node_mask: tf.Tensor
+
   # A tensor that contains feature vectors of nodes representing instructions in
   # the order in which they are in the basic block, i.e. in the same order
   # instructions appear in ModelBase._output_tensor_deltas.
@@ -237,12 +246,18 @@ def _create_graph_network_resources(self) -> None:
         shape=(None,),
         name=GraphBuilderModelBase.INSTRUCTION_NODE_MASK_TENSOR_NAME,
     )
+    self._context_node_mask = tf.placeholder(
+        dtype=tf.dtypes.bool,
+        shape=(None,),
+        name=GraphBuilderModelBase.CONTEXT_NODE_MASK_TENSOR_NAME,
+    )
 
   # @Override
   def _create_readout_network_resources(self) -> None:
     super()._create_readout_network_resources()
     self._instruction_features = tf.boolean_mask(
-        self._graphs_tuple_outputs.nodes, self._instruction_node_mask
+        self._graphs_tuple_outputs.nodes,
+        self._instruction_node_mask & ~self._context_node_mask,
     )
 
   # @Override
@@ -256,6 +271,9 @@ def _make_batch_feed_dict(self) -> model_base.FeedDict:
     feed_dict[self._instruction_node_mask] = np.array(
         self._batch_graph_builder.instruction_node_mask, dtype=bool
     )
+    feed_dict[self._context_node_mask] = np.array(
+        self._batch_graph_builder.context_node_mask, dtype=bool
+    )
     feed_dict[self._instruction_annotations] = (
         self._batch_graph_builder.instruction_annotations
     )
@@ -311,7 +329,10 @@ def _make_batch_graphs_tuple(self):
 
   # @Override
   def _add_basic_block_to_batch(self, block: basic_block.BasicBlock) -> None:
-    basic_block_was_added = self._batch_graph_builder.add_basic_block(block)
+    # Add context to the basic block graph only for seq2seq models.
+    basic_block_was_added = self._batch_graph_builder.add_basic_block(
+        block, add_context=self.use_deltas
+    )
     if not basic_block_was_added:
       # TODO(ondrasej): Better handling of blocks that can't be added to the
       # batch. For now, we just let the exception propagate out of the model and

gematria/granite/python/graph_builder_model_base_test.py

@@ -215,6 +215,39 @@ def test_train_seq2seq_model(self, loss_type, loss_normalization):
         model, self.blocks_with_throughput[0:1], num_epochs=50
     )
 
+  @parameterized.named_parameters(
+      *model_test.LOSS_TYPES_AND_LOSS_NORMALIZATIONS
+  )
+  def test_train_seq2seq_context_model(self, loss_type, loss_normalization):
+    blocks_with_throughput = self.blocks_with_throughput[0:1]
+    # Altered version of the regular basic blocks where all instructions belong
+    # to the context - i.e. an empty blocks with contexts matching the original
+    # basic blocks.
+    altered_blocks_with_throughput = [
+        throughput.BasicBlockWithThroughput(
+            block=basic_block.BasicBlock(
+                instructions=original.block.instructions,
+                back_context=original.block.instructions,
+                front_context=original.block.instructions,
+            ),
+            throughputs=original.throughputs,
+        )
+        for original in blocks_with_throughput
+    ]
+
+    model = TestGraphBuilderModel(
+        tokens=self.tokens,
+        loss_type=loss_type,
+        use_deltas=True,
+        use_delta_loss=False,
+        loss_normalization=loss_normalization,
+        num_message_passing_iterations=1,
+    )
+    model.initialize()
+    self.check_training_model(
+        model, altered_blocks_with_throughput, num_epochs=50
+    )
+
   def test_validate_basic_block(self):
     model = TestGraphBuilderModel(
         tokens=self.tokens, num_message_passing_iterations=1