class_docstr: One-parameter logistic item-response theory (IRT) model.
  
  This models a set of students answering a set of questions, and being scored
  whether they get the question correct or not. Each student is associated
  with a scalar `centered_student_ability`, and each question is associated
  with a scalar `question_difficulty`. Additionally, a scalar
  `mean_student_ability` is shared between all the students. This corresponds
  to the [1PL item-response theory][1] model:
  
  ```none
  mean_student_ability ∼ N(loc=0.75, scale=1)
  
  for i in range(num_students):
    centered_student_ability[i] ∼ N(loc=0, scale=1)
  
  for j in range(num_questions):
    question_difficulty[j] = N(loc=0, scale=1)
  
  for i, j in student_question_pairs:
    correct[i, j] = Bernoulli(logit=centered_student_ability[i] -
                              question_difficulty[j] + mean_student_ability)
  ```
  
  The data 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.
  
  #### References
  
  1. https://en.wikipedia.org/wiki/Item_response_theory
class_docstr_tok: ["One-parameter","logistic","item-response","theory","(","IRT",")","model",".","This","models","a","set","of","students","answering","a","set","of","questions",",","and","being","scored","whether","they","get","the","question","correct","or","not",".","Each","student","is","associated","with","a","scalar","`","centered_student_ability","`",",","and","each","question","is","associated","with","a","scalar","`","question_difficulty","`",".","Additionally",",","a","scalar","`","mean_student_ability","`","is","shared","between","all","the","students",".","This","corresponds","to","the","[","1PL","item-response","theory","]","[","1","]","model",":","``","`","none","mean_student_ability","\u223c","N","(","loc=0.75",",","scale=1",")","for","i","in","range","(","num_students",")",":","centered_student_ability","[","i","]","\u223c","N","(","loc=0",",","scale=1",")","for","j","in","range","(","num_questions",")",":","question_difficulty","[","j","]","=","N","(","loc=0",",","scale=1",")","for","i",",","j","in","student_question_pairs",":","correct","[","i",",","j","]","=","Bernoulli","(","logit=centered_student_ability","[","i","]","-","question_difficulty","[","j","]","+","mean_student_ability",")","``","`","The","data","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","did","n't",".","#","#","#","#","References","1.","https",":","\/\/en.wikipedia.org\/wiki\/Item_response_theory"]
docstr: Construct the item-response theory model.
  
  Args:
    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`. Can be `None`, in which
      case test-related sample transformations are not computed.
    test_question_ids: Integer `Tensor` with shape `[num_test_points]`.
      Testing question ids, ranging from 0 to `num_questions`. Can be `None`,
      in which case test-related sample transformations are not computed.
    test_correct: Integer `Tensor` with shape `[num_test_points]`. Whether the
      student in the testing set answered the question correctly, either 0 or
      1. Can be `None`, in which case test-related sample transformations are
      not computed.
    name: Python `str` name prefixed to Ops created by this class.
    pretty_name: A Python `str`. The pretty name of this model.
  
  Raises:
    ValueError: If `test_student_ids`, `test_question_ids` or `test_correct`
      are not either all `None` or are all specified.
    ValueError: If the parallel arrays are not all of the same size.
docstr_tok: ["Construct","the","item-response","theory","model",".","Args",":","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","`",".","Can","be","`","None","`",",","in","which","case","test-related","sample","transformations","are","not","computed",".","test_question_ids",":","Integer","`","Tensor","`","with","shape","`","[","num_test_points","]","`",".","Testing","question","ids",",","ranging","from","0","to","`","num_questions","`",".","Can","be","`","None","`",",","in","which","case","test-related","sample","transformations","are","not","computed",".","test_correct",":","Integer","`","Tensor","`","with","shape","`","[","num_test_points","]","`",".","Whether","the","student","in","the","testing","set","answered","the","question","correctly",",","either","0","or","1",".","Can","be","`","None","`",",","in","which","case","test-related","sample","transformations","are","not","computed",".","name",":","Python","`","str","`","name","prefixed","to","Ops","created","by","this","class",".","pretty_name",":","A","Python","`","str","`",".","The","pretty","name","of","this","model",".","Raises",":","ValueError",":","If","`","test_student_ids","`",",","`","test_question_ids","`","or","`","test_correct","`","are","not","either","all","`","None","`","or","are","all","specified",".","ValueError",":","If","the","parallel","arrays","are","not","all","of","the","same","size","."]
code: def __init__(
      self,
      train_student_ids,
      train_question_ids,
      train_correct,
      test_student_ids=None,
      test_question_ids=None,
      test_correct=None,
      name="item_response_theory",
      pretty_name="Item-Response Theory",
  ):
      """Construct the item-response theory model.
  
      Args:
        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`. Can be `None`, in which
          case test-related sample transformations are not computed.
        test_question_ids: Integer `Tensor` with shape `[num_test_points]`.
          Testing question ids, ranging from 0 to `num_questions`. Can be `None`,
          in which case test-related sample transformations are not computed.
        test_correct: Integer `Tensor` with shape `[num_test_points]`. Whether the
          student in the testing set answered the question correctly, either 0 or
          1. Can be `None`, in which case test-related sample transformations are
          not computed.
        name: Python `str` name prefixed to Ops created by this class.
        pretty_name: A Python `str`. The pretty name of this model.
  
      Raises:
        ValueError: If `test_student_ids`, `test_question_ids` or `test_correct`
          are not either all `None` or are all specified.
        ValueError: If the parallel arrays are not all of the same size.
      """
      with tf.name_scope(name):
          test_data_present = (
              e is not None
              for e in [
                  test_student_ids,
                  test_question_ids,
                  test_correct,
              ]
          )
          self._have_test = all(test_data_present)
          if not self._have_test and any(test_data_present):
              raise ValueError(
                  "`test_student_ids`, `test_question_ids` and "
                  "`test_correct` must either all be `None` or "
                  "all be specified. Got: test_student_ids={}, "
                  "test_question_ids={}, test_correct={}".format(
                      test_student_ids, test_question_ids, test_correct
                  )
              )
          if not (
              train_student_ids.shape[0]
              == train_question_ids.shape[0]
              == train_correct.shape[0]
          ):
              raise ValueError(
                  "`train_student_ids`, `train_question_ids` and "
                  "`train_correct` must all have the same length. "
                  "Got: {} {} {}".format(
                      train_student_ids.shape[0],
                      train_question_ids.shape[0],
                      train_correct.shape[0],
                  )
              )
  
          max_student_id = train_student_ids.max()
          max_question_id = train_question_ids.max()
          if self._have_test:
              max_student_id = max(
                  max_student_id, test_student_ids.max()
              )
              max_question_id = max(
                  max_question_id, test_question_ids.max()
              )
  
          self._num_students = max_student_id + 1
          self._num_questions = max_question_id + 1
  
          # TODO(siege): Make it an option to use a sparse encoding. The dense
          # encoding is only efficient when the dataset is not very sparse to begin
          # with.
          train_dense_y, train_y_mask = self._sparse_to_dense(
              train_student_ids,
              train_question_ids,
              train_correct,
          )
  
          self._prior_dist = tfd.JointDistributionNamed(
              dict(
                  mean_student_ability=tfd.Normal(0.75, 1.0),
                  centered_student_ability=tfd.Sample(
                      tfd.Normal(0.0, 1.0),
                      self._num_students,
                  ),
                  question_difficulty=tfd.Sample(
                      tfd.Normal(0.0, 1.0),
                      self._num_questions,
                  ),
              )
          )
  
          def observation_noise_fn(
              mean_student_ability,
              centered_student_ability,
              question_difficulty,
          ):
              """Creates the observation noise distribution."""
              logits = (
                  mean_student_ability[..., tf.newaxis, tf.newaxis]
                  + centered_student_ability[..., tf.newaxis]
                  - question_difficulty[..., tf.newaxis, :]
              )
              return tfd.Bernoulli(logits)
  
          self._observation_noise_fn = observation_noise_fn
  
          def log_likelihood_fn(
              dense_y, y_mask, reduce_sum=True, **params
          ):
              """The log_likelihood function."""
              log_likelihood = observation_noise_fn(**params).log_prob(
                  dense_y
              )
              log_likelihood = tf.where(
                  y_mask, log_likelihood, tf.zeros_like(log_likelihood)
              )
              if reduce_sum:
                  return tf.reduce_sum(log_likelihood, [-1, -2])
              else:
                  return log_likelihood
  
          self._train_log_likelihood_fn = functools.partial(
              log_likelihood_fn,
              dense_y=train_dense_y,
              y_mask=train_y_mask,
          )
  
          dtype = self._prior_dist.dtype
  
          sample_transformations = {
              "identity": model.Model.SampleTransformation(
                  fn=lambda params: params,
                  pretty_name="Identity",
                  dtype=dtype,
              )
          }
          if self._have_test:
              if not (
                  test_student_ids.shape[0]
                  == test_question_ids.shape[0]
                  == test_correct.shape[0]
              ):
                  raise ValueError(
                      "`test_student_ids`, `test_question_ids` and "
                      "`test_correct` must all have the same length. "
                      "Got: {} {} {}".format(
                          test_student_ids.shape[0],
                          test_question_ids.shape[0],
                          test_correct.shape[0],
                      )
                  )
              test_dense_y, test_y_mask = self._sparse_to_dense(
                  test_student_ids,
                  test_question_ids,
                  test_correct,
              )
              test_log_likelihood_fn = functools.partial(
                  log_likelihood_fn,
                  dense_y=test_dense_y,
                  y_mask=test_y_mask,
              )
  
              sample_transformations[
                  "test_nll"
              ] = model.Model.SampleTransformation(
                  fn=lambda params: test_log_likelihood_fn(**params),
                  pretty_name="Test NLL",
              )
  
              def _per_example_test_nll(params):
                  """Computes per-example test NLL."""
                  dense_nll = test_log_likelihood_fn(
                      reduce_sum=False, **params
                  )
                  return self._dense_to_sparse(
                      test_student_ids, test_question_ids, dense_nll
                  )
  
              sample_transformations[
                  "per_example_test_nll"
              ] = model.Model.SampleTransformation(
                  fn=_per_example_test_nll,
                  pretty_name="Per-example Test NLL",
              )
  
      self._train_student_ids = train_student_ids
      self._train_question_ids = train_question_ids
      self._test_student_ids = test_student_ids
      self._test_question_ids = test_question_ids
  
      super(ItemResponseTheory, self).__init__(
          default_event_space_bijector=tf.nest.map_structure(
              lambda _: tfb.Identity(), self._prior_dist.dtype
          ),
          event_shape=self._prior_dist.event_shape,
          dtype=dtype,
          name=name,
          pretty_name=pretty_name,
          sample_transformations=sample_transformations,
      )
  
file_path: probability/spinoffs/inference_gym/inference_gym/targets/item_response_theory.py
imports: import functools
  import numpy
  import tensorflow.compat.v2
  import tensorflow_probability
  from tensorflow_probability.python.internal import prefer_static
  from inference_gym.internal import data
  from inference_gym.targets import bayesian_model
  from inference_gym.targets import model
  from inference_gym.targets.ground_truth import synthetic_item_response_theory
label_desc: Use image node_id 1 to create a new ItemResponseTheory object from inherited base classes: bayesian_model with example: obj = ItemResponseTheory(train_student_ids, train_question_ids, train_correct, test_student_ids, test_question_ids, test_correct, name, pretty_name)