def test_generate_data_produces_examples_of_correct_shape(self):
    """Test examples have correct input and output shapes.

    Notes:

      * Loops over all AllenBrainImage2image* problems.

    """

    with TemporaryDirectory() as tmp_dir:
      mock_raw_data(tmp_dir, raw_dim=256, num_images=100)
      with TemporaryDirectory() as data_dir:
        for problem_obj in self.all_problems:
          problem_object = problem_obj()

          problem_object.generate_data(data_dir, tmp_dir)

          for mode in [Modes.TRAIN, Modes.EVAL]:

            dataset = problem_object.dataset(mode, data_dir)
            example = tfe.Iterator(dataset).next()

            num_channels = problem_object.num_channels

            # Check that the input tensor has the right shape
            input_dim = problem_object.input_dim
            self.assertEqual(example["inputs"].numpy().shape,
                             (input_dim, input_dim, num_channels))

            # Check that the targets tensor has the right shape
            output_dim = problem_object.output_dim
            self.assertEqual(example["targets"].numpy().shape,
                             (output_dim, output_dim, num_channels))
Test examples have correct input and output shapes.

    Notes:

      * Loops over all AllenBrainImage2image* problems.

    
tensor2tensor/data_generators/allen_brain_test.py