code: def test_masked_softmax(self) -> None:
      # Postive value case
      x = torch.tensor([[15.0, 6.0, 9.0], [3.0, 2.0, 1.0]])
      temperature = 1
      mask = torch.tensor([[1.0, 0.0, 1.0], [0.0, 1.0, 1.0]])
      # pyre-fixme[6]: For 2nd param expected `float` but got `Tensor`.
      out = masked_softmax(x, mask, temperature)
      expected_out = torch.tensor(
          [[0.9975, 0.0000, 0.0025], [0, 0.7311, 0.2689]]
      )
      npt.assert_array_almost_equal(out, expected_out, 4)
  
      # Postive value case (masked value goes to inf)
      x = torch.tensor([[150.0, 2.0]])
      temperature = 0.01
      mask = torch.tensor([[0.0, 1.0]])
      # pyre-fixme[6]: For 2nd param expected `float` but got `Tensor`.
      out = masked_softmax(x, mask, temperature)
      expected_out = torch.tensor([[0.0, 1.0]])
      npt.assert_array_almost_equal(out, expected_out, 4)
  
      # Negative value case
      x = torch.tensor([[-10.0, -1.0, -5.0]])
      temperature = 0.01
      mask = torch.tensor([[1.0, 1.0, 0.0]])
      # pyre-fixme[6]: For 2nd param expected `float` but got `Tensor`.
      out = masked_softmax(x, mask, temperature)
      expected_out = torch.tensor([[0.0, 1.0, 0.0]])
      npt.assert_array_almost_equal(out, expected_out, 4)
  
      # All values in a row are masked case
      x = torch.tensor([[-5.0, 4.0, 3.0], [2.0, 1.0, 2.0]])
      temperature = 1
      mask = torch.tensor([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]])
      # pyre-fixme[6]: For 2nd param expected `float` but got `Tensor`.
      out = masked_softmax(x, mask, temperature)
      expected_out = torch.tensor(
          [[0.0, 0.0, 0.0], [0.4223, 0.1554, 0.4223]]
      )
      npt.assert_array_almost_equal(out, expected_out, 4)
  
file_path: ReAgent/reagent/test/base/test_utils.py
imports: import unittest
  import numpy.testing
  import torch
  from reagent.core.torch_utils import masked_softmax, reorder_data_kjt, rescale_torch_tensor, shift_kjt_by_one, split_sequence_keyed_jagged_tensor
  from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
label_desc: Use image node_id 2 for calling the TestUtils obj's underlying member method code with example usage: obj.test_masked_softmax() without return types