code: def test_train_ds(self):
      expected_shape = [2, _TARGET_LENGTH]  # 2 devices.
      # For training we pack multiple short examples in one example.
      # *_position and *_segmentation indicate the boundaries.
      for batch in self.train_ds.take(3):
          self.assertEqual(
              {k: v.shape.as_list() for k, v in batch.items()},
              {
                  "inputs": expected_shape,
                  "inputs_position": expected_shape,
                  "inputs_segmentation": expected_shape,
                  "targets": expected_shape,
                  "targets_position": expected_shape,
                  "targets_segmentation": expected_shape,
              },
          )
  
file_path: flax/examples/wmt/input_pipeline_test.py
imports: import os
  import pathlib
  import tempfile
  from absl.testing import absltest
  import tensorflow_datasets
  from configs import default
  import input_pipeline
label_desc: Use image node_id 3 for calling the InputPipelineTest obj's underlying member method code with example usage: obj.test_train_ds() without return types