code: def test_source_length_mask():
      heads = 2
      batch_size = 4
      total_length = 6
      prepended_length = 3
      max_length = 8
      total_length_tensor = pt.tensor([total_length] * batch_size)
      # standard source-length mask array is [0, 0, 0, 0, 0, 0, 1, 1]
  
      for expand in [True, False]:
          for mask_prepended_tokens in [True, False]:
              if mask_prepended_tokens:
                  prepended_length_tensor = pt.tensor(
                      [prepended_length] * batch_size
                  )
                  # expected mask array is [1, 1, 1, 0, 0, 0, 1, 1]
                  masked_length = (
                      max_length - total_length + prepended_length
                  )
              else:
                  prepended_length_tensor = pt.tensor([0] * batch_size)
                  masked_length = max_length - total_length
              lengths = pt.stack(
                  (total_length_tensor, prepended_length_tensor), dim=1
              )
              mask = sockeye.layers.prepare_source_length_mask(
                  lengths,
                  heads,
                  max_length,
                  expand,
                  mask_prepended_tokens,
              )
              if expand:
                  total_size = batch_size * heads
                  assert mask.shape[1] == 1
              else:
                  total_size = batch_size
              assert mask.shape[0] == total_size
              assert pt.sum(mask) == total_size * masked_length
  
file_path: sockeye/test/unit/test_layers.py
imports: from math import pow, sqrt
  import numpy
  import pytest
  import torch
  import sockeye.constants
  import sockeye.layers
  import sockeye.transformer
  from sockeye.knn import KNNConfig, FaissIndexBuilder
label_desc: Use image node_id 2 for calling a global function with example usage: test_source_length_mask() without return types