code: def test_choice(generator_class):
      np_generator = {
          da.random.RandomState: np.random.RandomState,
          da.random.default_rng: np.random.default_rng,
      }
      np_dtype = (
          np_generator[generator_class]().choice(1, size=()).dtype
      )
      size = (10, 3)
      chunks = 4
      x = generator_class().choice(3, size=size, chunks=chunks)
      assert x.dtype == np_dtype
      assert x.shape == size
      res = x.compute()
      assert res.dtype == np_dtype
      assert res.shape == size
  
      py_a = [1, 3, 5, 7, 9]
      np_a = np.array(py_a, dtype="f8")
      da_a = da.from_array(np_a, chunks=2)
  
      for a in [py_a, np_a, da_a]:
          x = generator_class().choice(a, size=size, chunks=chunks)
          res = x.compute()
          expected_dtype = np.asarray(a).dtype
          assert x.dtype == expected_dtype
          assert res.dtype == expected_dtype
          assert set(np.unique(res)).issubset(np_a)
  
      np_p = np.array([0, 0.2, 0.2, 0.3, 0.3])
      da_p = da.from_array(np_p, chunks=2)
  
      for a, p in [(da_a, np_p), (np_a, da_p)]:
          x = generator_class().choice(a, size=size, chunks=chunks, p=p)
          res = x.compute()
          assert x.dtype == np_a.dtype
          assert res.dtype == np_a.dtype
          assert set(np.unique(res)).issubset(np_a[1:])
  
      np_dtype = (
          np_generator[generator_class]()
          .choice(1, size=(), p=np.array([1]))
          .dtype
      )
      x = generator_class().choice(5, size=size, chunks=chunks, p=np_p)
      res = x.compute()
      assert x.dtype == np_dtype
      assert res.dtype == np_dtype
  
      errs = [
          (-1, None),  # negative a
          (np_a[:, None], None),  # a must be 1D
          (np_a, np_p[:, None]),  # p must be 1D
          (np_a, np_p[:-2]),  # a and p must match
          (3, np_p),  # a and p must match
          (4, [0.2, 0.2, 0.3]),
      ]  # p must sum to 1
  
      for a, p in errs:
          with pytest.raises(ValueError):
              generator_class().choice(a, size=size, chunks=chunks, p=p)
  
      with pytest.raises(NotImplementedError):
          generator_class().choice(
              da_a, size=size, chunks=chunks, replace=False
          )
  
      # axis needs to be row based since a is 1D
      with pytest.raises(ValueError):
          da.random.default_rng().choice(
              da_a, size=size, chunks=chunks, axis=1
          )
  
      # Want to make sure replace=False works for a single-partition output array
      x = generator_class().choice(
          da_a, size=da_a.shape[0], chunks=-1, replace=False
      )
      res = x.compute()
      assert len(res) == len(np.unique(res))
  
imports: from __future__ import annotations
  import pytest
  import numpy
  import dask
  import dask.array
  from dask.array.core import Array
  from dask.array.utils import assert_eq
  from dask.multiprocessing import _dumps, _loads
  from dask.utils import key_split
label_desc: Use image node_id 24 for calling a global function with example usage: test_choice(generator_class) without return types