code: def testLayoutOptimizationConv2D(self):
      if not test_util.is_gpu_available():
          self.skipTest("No GPU available")
  
      # Compute convolution with input and filter of [1, 1, 1, 1] shape.
      # Verify that Grappler doesn't transpose Conv2D data format to NCHW.
      dataset = dataset_ops.Dataset.from_tensors((1, 1))
  
      def map_function(x, y):
          i = math_ops.cast(x, dtypes.float32)
          i = array_ops.reshape(i, [1, 1, 1, 1])
          f = math_ops.cast(y, dtypes.float32)
          f = array_ops.reshape(f, [1, 1, 1, 1])
          c = nn_ops.conv2d(i, f, strides=[1, 1, 1, 1], padding="VALID")
          return array_ops.reshape(c, ())
  
      dataset = dataset.map(map_function)
      self.assertDatasetProduces(dataset, expected_output=[1])
  
file_path: tensorflow-upstream/tensorflow/python/data/experimental/kernel_tests/optimization/grappler_test.py
imports: from absl.testing import parameterized
  from tensorflow.core.example import example_pb2
  from tensorflow.core.example import feature_pb2
  from tensorflow.python.data.kernel_tests import test_base
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.framework import combinations
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import test_util
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import cond
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import nn_ops
  from tensorflow.python.ops import parsing_ops
  from tensorflow.python.platform import test
label_desc: Use image node_id 2 for calling the GrapplerTest obj's underlying member method code with example usage: obj.testLayoutOptimizationConv2D() and returns: array_ops