code: def test_batch_node_features():
      (
          edge_features,
          edge_features_2,
          graphs,
          node_features,
          node_features_2,
      ) = generate_sequential_graphs(bs=5)
  
      batch = to_batch(graphs)
      batch_node_features = batch.batch_node_features
      batch_edge_features = batch.batch_edge_features
      assert torch.all(
          torch.eq(
              batch_node_features["node_feat"],
              torch.nn.utils.rnn.pad_sequence(
                  node_features, batch_first=True
              ),
          )
      )
      assert torch.all(
          torch.eq(
              batch_edge_features["edge_feat"],
              torch.nn.utils.rnn.pad_sequence(
                  edge_features, batch_first=True
              ),
          )
      )
  
      new_batch_edge_features = torch.stack(edge_features_2).unsqueeze(
          -1
      )
      batch.batch_edge_features["edge_feat"] = new_batch_edge_features
      graphs = from_batch(batch)
      for i in range(len(graphs)):
          g = graphs[i]
          assert torch.all(
              torch.eq(
                  g.edge_features["edge_feat"],
                  new_batch_edge_features[i],
              )
          )
  
imports: import gc
  import time
  import matplotlib.pyplot
  import torch
  import torch.nn
  import pytest
  from ...data.data import GraphData, from_batch, from_dgl, to_batch
  from ...data.utils import EdgeNotFoundException, SizeMismatchException
label_desc: Use image node_id 14 for calling a global function with example usage: test_batch_node_features() without return types