code: def test_space_sample_mask(space: Space, mask, n_trials: int = 100):
      """Tests that the sampling a space with a mask has the expected distribution.
  
      The implemented code is similar to the `test_space_sample` that considers the mask applied.
      """
      if isinstance(space, Box):
          # The box space can't have a sample mask
          assert mask is None
          return
      assert mask is not None
  
      space.seed(1)
      samples = np.array([space.sample(mask) for _ in range(n_trials)])
  
      if isinstance(space, Discrete):
          if np.any(mask == 1):
              expected_frequency = (
                  np.ones(space.n) * (n_trials / np.sum(mask)) * mask
              )
          else:
              expected_frequency = np.zeros(space.n)
              expected_frequency[0] = n_trials
          observed_frequency = np.zeros(space.n)
          for sample in samples:
              observed_frequency[sample - space.start] += 1
          degrees_of_freedom = max(np.sum(mask) - 1, 0)
  
          assert observed_frequency.shape == expected_frequency.shape
          assert np.sum(observed_frequency) == n_trials
          assert np.sum(expected_frequency) == n_trials
          variance = np.sum(
              np.square(expected_frequency - observed_frequency)
              / np.clip(expected_frequency, 1, None)
          )
          assert variance < CHI_SQUARED[degrees_of_freedom]
      elif isinstance(space, MultiBinary):
          expected_frequency = (
              np.ones(space.shape)
              * np.where(mask == 2, 0.5, mask)
              * n_trials
          )
          print(expected_frequency)
          observed_frequency = np.sum(samples, axis=0)
          assert (
              space.shape
              == expected_frequency.shape
              == observed_frequency.shape
          )
  
          variance = (
              2
              * np.square(observed_frequency - expected_frequency)
              / np.clip(expected_frequency, 1, None)
          )
          assert variance.shape == space.shape
          assert np.all(variance < CHI_SQUARED[1])
      elif isinstance(space, MultiDiscrete):
          # Due to the multi-axis capability of MultiDiscrete, these functions need to be recursive and that the expected / observed numpy are of non-regular shapes
          def _generate_frequency(
              _dim: Union[np.ndarray, int], _mask, func: callable
          ) -> List:
              if isinstance(_dim, np.ndarray):
                  return [
                      _generate_frequency(sub_dim, sub_mask, func)
                      for sub_dim, sub_mask in zip(_dim, _mask)
                  ]
              else:
                  return func(_dim, _mask)
  
          def _update_observed_frequency(obs_sample, obs_freq):
              if isinstance(obs_sample, np.ndarray):
                  for sub_sample, sub_freq in zip(obs_sample, obs_freq):
                      _update_observed_frequency(sub_sample, sub_freq)
              else:
                  obs_freq[obs_sample] += 1
  
          def _exp_freq_fn(_dim: int, _mask: np.ndarray):
              if np.any(_mask == 1):
                  assert _dim == len(_mask)
                  return (
                      np.ones(_dim) * (n_trials / np.sum(_mask)) * _mask
                  )
              else:
                  freq = np.zeros(_dim)
                  freq[0] = n_trials
                  return freq
  
          expected_frequency = _generate_frequency(
              space.nvec,
              mask,
              lambda dim, _mask: _exp_freq_fn(dim, _mask),
          )
          observed_frequency = _generate_frequency(
              space.nvec, mask, lambda dim, _: np.zeros(dim)
          )
          for sample in samples:
              _update_observed_frequency(sample, observed_frequency)
  
          def _chi_squared_test(dim, _mask, exp_freq, obs_freq):
              if isinstance(dim, np.ndarray):
                  for (
                      sub_dim,
                      sub_mask,
                      sub_exp_freq,
                      sub_obs_freq,
                  ) in zip(dim, _mask, exp_freq, obs_freq):
                      _chi_squared_test(
                          sub_dim, sub_mask, sub_exp_freq, sub_obs_freq
                      )
              else:
                  assert exp_freq.shape == (
                      dim,
                  ) and obs_freq.shape == (dim,)
                  assert np.sum(obs_freq) == n_trials
                  assert np.sum(exp_freq) == n_trials
                  _variance = np.sum(
                      np.square(exp_freq - obs_freq)
                      / np.clip(exp_freq, 1, None)
                  )
                  _degrees_of_freedom = max(np.sum(_mask) - 1, 0)
                  assert _variance < CHI_SQUARED[_degrees_of_freedom]
  
          _chi_squared_test(
              space.nvec, mask, expected_frequency, observed_frequency
          )
      elif isinstance(space, Text):
          length, charlist_mask = mask
  
          if length is None:
              expected_length = (
                  space.min_length
                  + (space.max_length - space.min_length) / 2
              )
          else:
              expected_length = length
  
          if np.any(charlist_mask == 1):
              expected_frequency = (
                  np.ones(len(space.character_set))
                  * n_trials
                  * expected_length
                  / np.sum(charlist_mask)
                  * charlist_mask
              )
          else:
              expected_frequency = np.zeros(len(space.character_set))
  
          observed_frequency = np.zeros(len(space.character_set))
          for sample in samples:
              for char in sample:
                  observed_frequency[space.character_index(char)] += 1
  
          degrees_of_freedom = max(np.sum(charlist_mask) - 1, 0)
  
          assert observed_frequency.shape == expected_frequency.shape
          assert np.sum(observed_frequency) == sum(
              len(sample) for sample in samples
          )
  
          variance = np.sum(
              np.square(expected_frequency - observed_frequency)
              / np.clip(expected_frequency, 1, None)
          )
          if degrees_of_freedom == 26:
              # scipy.stats.chi2.isf(0.05, df=29)
              assert variance < 38.88513865983007
          elif degrees_of_freedom == 31:
              # scipy.stats.chi2.isf(0.05, df=31)
              assert variance < 44.985343280365136
          else:
              assert variance < CHI_SQUARED[degrees_of_freedom]
      else:
          raise NotImplementedError()
  
imports: import copy
  import itertools
  import json
  import pickle
  import tempfile
  from typing import List, Union
  import numpy
  import pytest
  from gym.spaces import Box, Discrete, MultiBinary, MultiDiscrete, Space, Text
  from gym.utils import seeding
  from gym.utils.env_checker import data_equivalence
  from tests.spaces.utils import TESTING_FUNDAMENTAL_SPACES, TESTING_FUNDAMENTAL_SPACES_IDS, TESTING_SPACES, TESTING_SPACES_IDS
label_desc: Use image node_id 4 for calling a global function with example usage: test_space_sample_mask(space, mask, n_trials) and returns: unknown, func, unknown, freq