// RUN: tf-opt %s -allow-unregistered-dialect --tf-strip-tf-attributes --split-input-file | FileCheck %s

// CHECK-LABEL: strips_attributes
// CHECK-NOT: tf
func.func @strips_attributes(%arg0: tensor<32x28x28x1xf32> {tf._user_specified_name = "x"},
                             %arg1: tensor<3x3x1x5xf32> {tf._user_specified_name = "w1"},
                             %arg2: tensor<5xf32> {tf._user_specified_name = "b1"},

  %7 = "mhlo.broadcast_in_dim"(%3) <{broadcast_dimensions = dense<3> : tensor<1xi64>}> : (tensor<5xf32>) -> tensor<32x28x28x5xf32>
  %8 = mhlo.add %6, %7 : tensor<32x28x28x5xf32>
  %9 = mhlo.maximum %8, %1 : tensor<32x28x28x5xf32>
  %10 = "mhlo.reshape"(%9) : (tensor<32x28x28x5xf32>) -> tensor<32x3920xf32>

// 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> {

  %1 = "tf.Const"() {value = dense<1.000000e+00> : tensor<1xf32>} : () -> tensor<1xf32>
  %2 = "tf.Sub"(%0, %1): (tensor<15x28x28x1xf32>, tensor<1xf32>) -> tensor<15x28x28x1xf32>
  func.return %2 : tensor<15x28x28x1xf32>
// CHECK: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor<1xf32>}> : () -> tensor<1xf32>

// CHECK: return %[[ADD]] : tensor<1x1x64xf32>
// CHECK: }

// -----

// CHECK-LABEL: @conv_with_bias_dynamic_fn(
// CHECK-SAME:                    %[[ARG_0:.*]]: tensor<?x28x28x1xf32>
func.func @conv_with_bias_dynamic_fn(%arg0: tensor<?x28x28x1xf32>) -> tensor<?x28x28x16xf32> {
  %0 = stablehlo.constant dense<2.000000e+00> : tensor<3x3x1x16xf32>
  %1 = stablehlo.constant dense<2.000000e+00> : tensor<16xf32>

  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 = "tf.Add"(%0, %1): (tensor<15x28x28x1xf32>, tensor<15x28x28x1xf32>) -> tensor<15x28x28x1xf32>
  func.return %2 : tensor<15x28x28x1xf32>

  // CHECK-LABEL: conv2d_backprop_input
  // CHECK: %[[CST:.*]] = arith.constant dense<[2, 0, 1, 3]> : tensor<4xi32>