// CHECK: return %[[ADDV2_0]] : tensor<15x28x28x1xf32>
}

// CHECK-LABEL: conv2d_backprop_input_with_sub
func.func @conv2d_backprop_input_with_sub(%arg0: tensor<4xi32>, %arg1: tensor<3x3x1x32xf32>, %arg2: tensor<15x14x14x32xf32>) -> tensor<15x28x28x1xf32> {

// generated GraphDef to make sure the export is valid.

module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 890 : i32}}  {
  func.func @main(%arg0: tensor<100x28x28x3xf32> {}) -> (tensor<*xf32> {}) attributes {tf.entry_function = {control_outputs = "", inputs = "a:0", outputs = "conv2d_2/BiasAdd:0"}} {
    %0 = tf_executor.graph {

// CHECK-LABEL: convolution_3d
func.func @convolution_3d(%arg0: tensor<1x4x28x28x1xf32>, %arg1: tensor<2x3x3x1x16xf32>) -> tensor<1x3x26x26x16xf32> {
  %0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, 2, f]x[0, 1, 2, i, o]->[b, 0, 1, 2, f], window = {} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x4x28x28x1xf32>, tensor<2x3x3x1x16xf32>) -> tensor<1x3x26x26x16xf32>
  return %0 : tensor<1x3x26x26x16xf32>
}

  return model


class TestModule(tf.Module):

  def __init__(self):
    super(TestModule, self).__init__()
    self.model = mnist_model()

  # CHECK: func {{@[a-zA-Z_0-9]+}}(%arg0: tensor<1x28x28x1xf32> {tf._user_specified_name = "x", tf_saved_model.index_path = [0]}
  # CHECK: attributes {{.*}} tf_saved_model.exported_names = ["my_predict"]
  @tf.function(input_signature=[
      tf.TensorSpec([1, 28, 28, 1], tf.float32),

  %2:2 = "tf.IteratorGetNext"(%0) {_class = ["loc:@iterator"], device = ""} : (tensor<!tf_type.resource>) -> (tensor<200x28x28x1xf32>, tensor<200x10xf32>)
  // CHECK: tf.Iterator
  // CHECK: tf.MakeIterator
  // CHECK: tf.IteratorGetNext
  return