code: def synthetic_item_response_theory(
      train_fraction=1.0,
      shuffle=True,
      shuffle_seed=1337,
  ):
      """Synthetic dataset sampled from the ItemResponseTheory model.
  
      This dataset is a simulation of 400 students each answering a subset of
      100 unique questions, with a total of 30012 questions answered.
  
      The dataset is split into train and test portions by randomly partitioning the
      student-question-response triplets. This has two consequences. First, the
      student and question ids are shared between test and train sets. Second, there
      is a possibility of some students or questions not being present in both sets.
  
      The train set will contain `num_train_points = int(train_fraction * 30012)`
      triples and test set will contain
      `num_test_points = 30012 - num_train_points` triplets.
  
      The triples are encoded into three parallel arrays per set. I.e.
      `*_correct[i]  == 1` means that student `*_student_ids[i]` answered question
      `*_question_ids[i]` correctly; `*_correct[i] == 0` means they didn't.
  
      Args:
        train_fraction: What fraction of the data to put in the training set.
        shuffle: Whether to shuffle the dataset.
        shuffle_seed: Seed to use when shuffling.
  
      Returns:
        dataset: A Dict with the following keys:
          `train_student_ids`: Integer `Tensor` with shape `[num_train_points]`.
            training student ids, ranging from 0 to `num_students`.
          `train_question_ids`: Integer `Tensor` with shape `[num_train_points]`.
            training question ids, ranging from 0 to `num_questions`.
          `train_correct`: Integer `Tensor` with shape `[num_train_points]`.
            Whether the student in the training set answered the question correctly,
            either 0 or 1.
          `test_student_ids`: Integer `Tensor` with shape `[num_test_points]`.
            Testing student ids, ranging from 0 to `num_students`.
          `test_question_ids`: Integer `Tensor` with shape `[num_test_points]`.
            Testing question ids, ranging from 0 to `num_questions`.
          `test_correct`: Integer `Tensor` with shape `[num_test_points]`.
            Whether the student in the testing set answered the question correctly,
            either 0 or 1.
      """
      student_ids = synthetic_item_response_theory_lib.STUDENT_IDS
      question_ids = synthetic_item_response_theory_lib.QUESTION_IDS
      correct = synthetic_item_response_theory_lib.CORRECT
  
      if shuffle:
          shuffle_idxs = np.arange(student_ids.shape[0])
          np.random.RandomState(shuffle_seed).shuffle(shuffle_idxs)
          student_ids = student_ids[shuffle_idxs]
          question_ids = question_ids[shuffle_idxs]
          correct = correct[shuffle_idxs]
  
      num_train = int(student_ids.shape[0] * train_fraction)
  
      return dict(
          train_student_ids=student_ids[:num_train],
          train_question_ids=question_ids[:num_train],
          train_correct=correct[:num_train],
          test_student_ids=student_ids[num_train:],
          test_question_ids=question_ids[num_train:],
          test_correct=correct[num_train:],
      )
  
imports: import numpy
  import tensorflow.compat.v2
  import tensorflow_probability
  from inference_gym.internal.datasets import brownian_motion_missing_middle_observations
  from inference_gym.internal.datasets import convection_lorenz_bridge
  from inference_gym.internal.datasets import sp500_closing_prices
  from inference_gym.internal.datasets import synthetic_item_response_theory
  from inference_gym.internal.datasets import synthetic_log_gaussian_cox_process
  from inference_gym.internal.datasets import synthetic_plasma_spectroscopy
  from inference_gym.internal.datasets import synthetic_plasma_spectroscopy_with_bump
label_desc: Use image node_id 12 for calling a global function with example usage: synthetic_item_response_theory(train_fraction, shuffle, shuffle_seed) and returns: dict