def testRetraceHatesBadlyRankedInputs(self):
    """Ensure Retrace notices inputs with the wrong rank."""
    # No problems if we create a Retrace using correctly-ranked arguments.
    proper_args = [self.lambda_, self.qs, self.targnet_qs, self.actions,
                   self.rewards, self.pcontinues, self.target_policy_probs,
                   self.behaviour_policy_probs]
    retrace_ops.retrace(*proper_args)

    # Now make a local copy of the args and try modifying each element to have
    # an inappropriate rank. We should get an error each time.
    for i in xrange(len(proper_args)):
      bad_args = list(proper_args)
      bad_args[i] = [bad_args[i]]
      with self.assertRaises(ValueError):
        retrace_ops.retrace(*bad_args)
Ensure Retrace notices inputs with the wrong rank.
trfl/retrace_ops_test.py