def testHigherRankInputsForIW(self):
    """Checks support for additional dimensions in inputs."""
    placeholders = {
        'log_rhos': tf.placeholder(dtype=tf.float32, shape=[None, None, 1]),
        'discounts': tf.placeholder(dtype=tf.float32, shape=[None, None, 1]),
        'rewards': tf.placeholder(dtype=tf.float32, shape=[None, None, 42]),
        'values': tf.placeholder(dtype=tf.float32, shape=[None, None, 42]),
        'bootstrap_value': tf.placeholder(dtype=tf.float32, shape=[None, 42])
    }
    output = vtrace_ops.vtrace_from_importance_weights(**placeholders)
    self.assertEqual(output.vs.shape.as_list()[-1], 42)
Checks support for additional dimensions in inputs.
trfl/vtrace_ops_test.py