def test_output_values_larger_rank(self, broadcaster, expected):
    """Test the broadcasted output value."""
    input_graph = utils_tf.data_dicts_to_graphs_tuple(
        [SMALL_GRAPH_1, SMALL_GRAPH_2])
    input_graph = input_graph.map(
        lambda v: tf.reshape(v, [v.get_shape().as_list()[0]] + [2, -1]))
    broadcasted = broadcaster(input_graph)
    with tf.Session() as sess:
      broadcasted_out = sess.run(broadcasted)
    self.assertNDArrayNear(
        np.reshape(np.array(expected, dtype=np.float32),
                   [len(expected)] + [2, -1]),
        broadcasted_out,
        err=1e-4)