code: def test_multinomial(self):
      self.assertRaises(
          AssertionError,
          lambda: Tensor(2).multinomial(1, replacement=False),
      )
      self.assertRaises(
          AssertionError,
          lambda: Tensor([1, 9]).multinomial(0, replacement=False),
      )
  
      def _check_with_torch(w, num_samples, replacement):
          tiny_res = Tensor(w).multinomial(
              num_samples, replacement=replacement
          )
          torch_res = torch.tensor(w).multinomial(
              num_samples, replacement=replacement
          )
          self.assertEqual(tiny_res.shape, torch_res.shape)
          if torch_res.ndim == 1:
              tiny_res = tiny_res.unsqueeze(0)
              torch_res = torch_res.unsqueeze(0)
          for i in range(torch_res.shape[0]):
              self.assertTrue(
                  equal_distribution(
                      lambda *_: tiny_res[i], lambda _: torch_res[i]
                  )
              )
  
      _check_with_torch(
          w=[0.231, 0.0, 1.0, 0.5], num_samples=2000, replacement=True
      )
      _check_with_torch(
          w=[[0.2, 0.8]], num_samples=2000, replacement=True
      )  # 2D but only 1 row
      _check_with_torch(
          w=[[0.453, 0.0, 1.0, 0.81], [0.1, 0.8, 0.0, 0.1]],
          num_samples=2000,
          replacement=True,
      )
      # no-replacement isn't supported, unless taking only one sample
      w = [0.1, 0.9]
      self.assertRaises(
          AssertionError,
          lambda: Tensor(w).multinomial(100, replacement=False),
      )
      tiny_samples = [
          Tensor(w).multinomial(1, replacement=False).numpy().item()
          for _ in range(1000)
      ]
      torch_samples = [
          torch.tensor(w).multinomial(1, replacement=False).item()
          for _ in range(1000)
      ]
      self.assertTrue(
          equal_distribution(
              lambda *_: Tensor(tiny_samples),
              lambda _: torch.tensor(torch_samples),
          )
      )
  
imports: import math
  import unittest
  import numpy
  import torch
  from tinygrad import nn, dtypes, Tensor
  from functools import partial
label_desc: Use image node_id 10 for calling the TestRandomness obj's underlying member method code with example usage: obj.test_multinomial() without return types