code: def test_dataset_with_ngrams(synthetic_dataset, reader_factory):
      ngram = NGram(
          {
              -1: [TestSchema.id, TestSchema.image_png],
              2: [TestSchema.id],
          },
          100,
          TestSchema.id,
      )
      with reader_factory(
          synthetic_dataset.url, schema_fields=ngram, num_epochs=1
      ) as reader:
          dataset = make_petastorm_dataset(reader)
          next_sample = dataset.make_one_shot_iterator().get_next()
          with tf.Session() as sess:
              while True:
                  try:
                      actual = sess.run(next_sample)
                      assert actual[-1].id + 3 == actual[2].id
                      assert np.all(
                          actual[-1].image_png.shape > (10, 10, 3)
                      )
                  except tf.errors.OutOfRangeError:
                      break
  
file_path: petastorm/petastorm/tests/test_tf_dataset.py
imports: import functools
  import operator
  from copy import copy
  import numpy
  import pytest
  import tensorflow.compat.v1
  from petastorm import make_reader, make_batch_reader
  from petastorm.ngram import NGram
  from petastorm.predicates import in_lambda
  from petastorm.tests.test_common import TestSchema
  from petastorm.tests.test_tf_utils import create_tf_graph
  from petastorm.tf_utils import make_petastorm_dataset
label_desc: Use image node_id 6 for calling a global function with example usage: test_dataset_with_ngrams(synthetic_dataset, reader_factory) without return types