code: def test_adjust_first_step_masking():
      first_step_mask = pt.tensor(
          [
              [0.0],
              [np.inf],
              [np.inf],
              [np.inf],
              [0.0],
              [np.inf],
              [np.inf],
              [np.inf],
          ]
      )
      target_prefix = pt.tensor([[1, 2], [1, 0]])
      adjust_first_step_mask = pt.tensor(
          [
              [0.0, 0.0, 0.0],
              [np.inf, np.inf, np.inf],
              [np.inf, np.inf, np.inf],
              [np.inf, np.inf, np.inf],
              [0.0, 0.0, 0.0],
              [np.inf, np.inf, 0.0],
              [np.inf, np.inf, 0.0],
              [np.inf, np.inf, 0.0],
          ]
      )
      assert (
          pt.equal(
              adjust_first_step_mask,
              utils.adjust_first_step_masking(
                  target_prefix, first_step_mask
              ),
          )
          == True
      )
      target_prefix = pt.tensor([[1, 0], [2, 3]])
      adjust_first_step_mask = pt.tensor(
          [
              [0.0, 0.0, 0.0],
              [np.inf, np.inf, 0.0],
              [np.inf, np.inf, 0.0],
              [np.inf, np.inf, 0.0],
              [0.0, 0.0, 0.0],
              [np.inf, np.inf, np.inf],
              [np.inf, np.inf, np.inf],
              [np.inf, np.inf, np.inf],
          ]
      )
      assert (
          pt.equal(
              adjust_first_step_mask,
              utils.adjust_first_step_masking(
                  target_prefix, first_step_mask
              ),
          )
          == True
      )
      target_prefix = pt.tensor([[1, 0, 0], [2, 3, 4]])
      adjust_first_step_mask = pt.tensor(
          [
              [0.0, 0.0, 0.0, 0.0],
              [np.inf, np.inf, 0.0, 0.0],
              [np.inf, np.inf, 0.0, 0.0],
              [np.inf, np.inf, 0.0, 0.0],
              [0.0, 0.0, 0.0, 0.0],
              [np.inf, np.inf, np.inf, np.inf],
              [np.inf, np.inf, np.inf, np.inf],
              [np.inf, np.inf, np.inf, np.inf],
          ]
      )
      assert (
          pt.equal(
              adjust_first_step_mask,
              utils.adjust_first_step_masking(
                  target_prefix, first_step_mask
              ),
          )
          == True
      )
      target_prefix = pt.tensor([[1, 0, 0, 0], [2, 3, 4, 5]])
      adjust_first_step_mask = pt.tensor(
          [
              [0.0, 0.0, 0.0, 0.0, 0.0],
              [np.inf, np.inf, 0.0, 0.0, 0.0],
              [np.inf, np.inf, 0.0, 0.0, 0.0],
              [np.inf, np.inf, 0.0, 0.0, 0.0],
              [0.0, 0.0, 0.0, 0.0, 0.0],
              [np.inf, np.inf, np.inf, np.inf, np.inf],
              [np.inf, np.inf, np.inf, np.inf, np.inf],
              [np.inf, np.inf, np.inf, np.inf, np.inf],
          ]
      )
      assert (
          pt.equal(
              adjust_first_step_mask,
              utils.adjust_first_step_masking(
                  target_prefix, first_step_mask
              ),
          )
          == True
      )
      first_step_mask = pt.tensor(
          [
              [0.0],
              [np.inf],
              [np.inf],
              [np.inf],
              [0.0],
              [np.inf],
              [np.inf],
              [np.inf],
              [0.0],
              [np.inf],
              [np.inf],
              [np.inf],
          ]
      )
      target_prefix = pt.tensor(
          [[1, 0, 0, 0], [1, 0, 0, 0], [2, 3, 4, 5]]
      )
      adjust_first_step_mask = pt.tensor(
          [
              [0.0, 0.0, 0.0, 0.0, 0.0],
              [np.inf, np.inf, 0.0, 0.0, 0.0],
              [np.inf, np.inf, 0.0, 0.0, 0.0],
              [np.inf, np.inf, 0.0, 0.0, 0.0],
              [0.0, 0.0, 0.0, 0.0, 0.0],
              [np.inf, np.inf, 0.0, 0.0, 0.0],
              [np.inf, np.inf, 0.0, 0.0, 0.0],
              [np.inf, np.inf, 0.0, 0.0, 0.0],
              [0.0, 0.0, 0.0, 0.0, 0.0],
              [np.inf, np.inf, np.inf, np.inf, np.inf],
              [np.inf, np.inf, np.inf, np.inf, np.inf],
              [np.inf, np.inf, np.inf, np.inf, np.inf],
          ]
      )
      assert (
          pt.equal(
              adjust_first_step_mask,
              utils.adjust_first_step_masking(
                  target_prefix, first_step_mask
              ),
          )
          == True
      )
      target_prefix = pt.tensor(
          [[1, 0, 0, 0], [1, 3, 0, 0], [2, 3, 4, 5]]
      )
      adjust_first_step_mask = pt.tensor(
          [
              [0.0, 0.0, 0.0, 0.0, 0.0],
              [np.inf, np.inf, 0.0, 0.0, 0.0],
              [np.inf, np.inf, 0.0, 0.0, 0.0],
              [np.inf, np.inf, 0.0, 0.0, 0.0],
              [0.0, 0.0, 0.0, 0.0, 0.0],
              [np.inf, np.inf, np.inf, 0.0, 0.0],
              [np.inf, np.inf, np.inf, 0.0, 0.0],
              [np.inf, np.inf, np.inf, 0.0, 0.0],
              [0.0, 0.0, 0.0, 0.0, 0.0],
              [np.inf, np.inf, np.inf, np.inf, np.inf],
              [np.inf, np.inf, np.inf, np.inf, np.inf],
              [np.inf, np.inf, np.inf, np.inf, np.inf],
          ]
      )
      assert (
          pt.equal(
              adjust_first_step_mask,
              utils.adjust_first_step_masking(
                  target_prefix, first_step_mask
              ),
          )
          == True
      )
      target_prefix = pt.tensor(
          [[0, 0, 0, 0], [1, 3, 0, 0], [2, 3, 4, 5]]
      )
      adjust_first_step_mask = pt.tensor(
          [
              [0.0, 0.0, 0.0, 0.0, 0.0],
              [np.inf, 0.0, 0.0, 0.0, 0.0],
              [np.inf, 0.0, 0.0, 0.0, 0.0],
              [np.inf, 0.0, 0.0, 0.0, 0.0],
              [0.0, 0.0, 0.0, 0.0, 0.0],
              [np.inf, np.inf, np.inf, 0.0, 0.0],
              [np.inf, np.inf, np.inf, 0.0, 0.0],
              [np.inf, np.inf, np.inf, 0.0, 0.0],
              [0.0, 0.0, 0.0, 0.0, 0.0],
              [np.inf, np.inf, np.inf, np.inf, np.inf],
              [np.inf, np.inf, np.inf, np.inf, np.inf],
              [np.inf, np.inf, np.inf, np.inf, np.inf],
          ]
      )
      assert (
          pt.equal(
              adjust_first_step_mask,
              utils.adjust_first_step_masking(
                  target_prefix, first_step_mask
              ),
          )
          == True
      )
      target_prefix = pt.tensor(
          [[0, 0, 0, 0], [0, 0, 0, 0], [2, 3, 4, 5]]
      )
      adjust_first_step_mask = pt.tensor(
          [
              [0.0, 0.0, 0.0, 0.0, 0.0],
              [np.inf, 0.0, 0.0, 0.0, 0.0],
              [np.inf, 0.0, 0.0, 0.0, 0.0],
              [np.inf, 0.0, 0.0, 0.0, 0.0],
              [0.0, 0.0, 0.0, 0.0, 0.0],
              [np.inf, 0.0, 0.0, 0.0, 0.0],
              [np.inf, 0.0, 0.0, 0.0, 0.0],
              [np.inf, 0.0, 0.0, 0.0, 0.0],
              [0.0, 0.0, 0.0, 0.0, 0.0],
              [np.inf, np.inf, np.inf, np.inf, np.inf],
              [np.inf, np.inf, np.inf, np.inf, np.inf],
              [np.inf, np.inf, np.inf, np.inf, np.inf],
          ]
      )
      assert (
          pt.equal(
              adjust_first_step_mask,
              utils.adjust_first_step_masking(
                  target_prefix, first_step_mask
              ),
          )
          == True
      )
  
imports: import gzip
  import math
  import os
  import re
  from tempfile import TemporaryDirectory
  import unittest
  import numpy
  import pytest
  import torch
  from sockeye import __version__
  from sockeye import constants
  from sockeye import utils
label_desc: Use image node_id 21 for calling a global function with example usage: test_adjust_first_step_masking() without return types