%0 = "tf.Conv2DBackpropInput"(%arg0, %arg1, %arg2) {strides = [1, 2, 2, 1], padding="SAME", dilations=[1, 1, 1, 1]}: (tensor<4xi32>, tensor<3x3x1x32xf32>, 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.func private @"map/while_body"(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<!tf_type.variant<tensor<*xf32>>>, %arg3: tensor<i32>, %arg4: tensor<!tf_type.resource<tensor<3x1xf32>>>, %arg5: tensor<?x3xf32>, %arg6: tensor<?x4xf32>) -> (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>, tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<?x3xf32>, tensor<?x4xf32>) {

    %8 = stablehlo.convert %7 : (tensor<1x4x2xi8>) -> tensor<1x4x2xf32>
    %9 = stablehlo.convert %2 : (tensor<2x3xi8>) -> tensor<2x3xf32>
    %10 = stablehlo.dot_general %8, %9, contracting_dims = [2] x [0] : (tensor<1x4x2xf32>, tensor<2x3xf32>) -> tensor<1x4x3xf32>
    %11 = stablehlo.convert %3 : (tensor<1x1x3xi32>) -> tensor<1x1x3xf32>

    // CHECK: [[GRAD_SQUARE:%.*]] = "tf.Mul"([[GRAD]], [[GRAD]]) : (tensor<f32>, tensor<f32>) -> tensor<f32>
    // CHECK: [[NEW_ACC:%.*]] = "tf.AddV2"([[OLD_ACC]], [[GRAD_SQUARE]]) : (tensor<*xf32>, tensor<f32>) -> tensor<*xf32>
    // CHECK: [[LR_MULTIPLY:%.*]] = "tf.Mul"([[LR]], [[GRAD]]) : (tensor<f32>, tensor<f32>) -> tensor<f32>

// CHECK-SAME: -> (tensor<*xi32> {mhlo.sharding = "\0A\0B\0C"}, tensor<*xi1> {mhlo.sharding = "\0D\0E\0F"})
func.func @func_with_sharding_after_read_variable(%arg0: tensor<*x!tf_type.resource<tensor<32xf32>>>, %arg1: tensor<*x!tf_type.resource<tensor<32xf32>>>) -> (tensor<*xi32>, tensor<*xi1>) {
  %0 = "tf.ReadVariableOp"(%arg0) : (tensor<*x!tf_type.resource<tensor<32xf32>>>) -> tensor<32xf32>

// CHECK: return %[[BIASADD_0]] : tensor<1x10x10x32xf32>
}

// CHECK-LABEL: add
func.func @add(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<1xf32> {
  %0 = "tf.AddV2"(%arg0, %arg1) : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
  func.return %0: tensor<1xf32>
// CHECK: %[[ADD_0:.*]] = "tf.AddV2"(%arg0, %arg1) : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
// CHECK: return %[[ADD_0]] : tensor<1xf32>
}

// CHECK-LABEL: softmax

    %0 = func.call @test_sdpa(%arg2, %arg3, %arg4, %arg5, %arg6) : (tensor<1x100x32x4xf32>,  tensor<1x500x4x4xf32>, tensor<1x500x4x4xf32>, tensor<1x1x100x500xf32>, tensor<f32>) -> tensor<1x100x32x4xf32>
    return %0: tensor<1x100x32x4xf32>
  }

  "tf.AssignVariableOp"(%arg1, %0#1) : (tensor<*x!tf_type.resource<tensor<i32>>>, tensor<i32>) -> ()
  func.return
}

// CHECK-LABEL: func @multiple_write_func
// CHECK-SAME: ({{%.*}}: tensor<i32>) -> (tensor<i32>, tensor<i32>)
func.func @multiple_write_func(%arg0: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
  func.return %arg0, %arg0 : tensor<i32>, tensor<i32>
}

  return
}
func.func private @false(%arg0: tensor<!tf_type.variant<tensor<f32>>>) -> tensor<!tf_type.variant<tensor<f32>>> {
  return %arg0 : tensor<!tf_type.variant<tensor<f32>>>
}
func.func private @true(%arg0: tensor<!tf_type.variant<tensor<f32>>>) -> tensor<!tf_type.variant<tensor<f32>>> {
  return %arg0 : tensor<!tf_type.variant<tensor<f32>>>
}

// -----

  func.return %0#2, %0#3, %0#4 : tensor<i32>, tensor<256x256xf32>, tensor<?x256x256xf32>
}