code: def test_safe_multinomial(self):
      mask = [
          [1, 0, 0, 0, 0],
          [1, 1, 0, 0, 0],
          [1, 1, 1, 0, 0],
          [1, 1, 1, 1, 0],
      ]
      tmask = torch.tensor(mask, dtype=torch.float32)
  
      for _ in range(5):
          random_scalar = torch.rand(1)
          samples = _safe_multinomial(tmask * random_scalar, 3)
          self.assertTupleEqual(samples.shape, (4, 3))
  
          # samples[0] is exactly determined
          self.assertConstant(samples[0], 0)
  
          self.assertGreaterEqual(samples[1].min(), 0)
          self.assertLessEqual(samples[1].max(), 1)
  
          # samples[2] is exactly determined
          self.assertSetEqual(set(samples[2].tolist()), {0, 1, 2})
  
          # samples[3] has enough sources, so must contain 3 distinct values.
          self.assertLessEqual(samples[3].max(), 3)
          self.assertEqual(len(set(samples[3].tolist())), 3)
  
imports: import unittest
  from typing import Callable
  import torch
  from pytorch3d.common.compat import meshgrid_ij
  from pytorch3d.ops import eyes
  from pytorch3d.renderer import MonteCarloRaysampler, MultinomialRaysampler, NDCGridRaysampler, NDCMultinomialRaysampler
  from pytorch3d.renderer.cameras import FoVOrthographicCameras, FoVPerspectiveCameras, OrthographicCameras, PerspectiveCameras
  from pytorch3d.renderer.implicit.raysampling import _jiggle_within_stratas, _safe_multinomial
  from pytorch3d.renderer.implicit.utils import ray_bundle_to_ray_points, ray_bundle_variables_to_ray_points
  from pytorch3d.transforms import Rotate
  from .common_testing import TestCaseMixin
  from .test_cameras import init_random_cameras
label_desc: Use image node_id 11 for calling the TestRaysampling obj's underlying member method code with example usage: obj.test_safe_multinomial() without return types