%0 = "tf.SpaceToBatchND"(%input, %block_shape, %paddings) : (tensor<3x5x7x10xf32>, tensor<2xi64>, tensor<2x2xi64>) -> tensor<36x2x3x10xf32>
  func.return %0 : tensor<36x2x3x10xf32>
}

// -----

// Test invalid tf.SpaceToBatchND

// CHECK-DAG: %[[cst_1:.*]] = "tf.Const"() <{value = dense<-1> : tensor<1xi64>}> : () -> tensor<1xi64>
// CHECK: %[[arg1_i64:.*]] = "tf.Cast"(%arg1) <{Truncate = false}> : (tensor<1xi32>) -> tensor<1xi64>
// CHECK: %[[tensor_scatter_update:.*]] = "tf.TensorScatterUpdate"(%[[cst]], %[[cst_0]], %[[arg1_i64]]) : (tensor<2xi64>, tensor<1x1xi64>, tensor<1xi64>) -> tensor<2xi64>

func.func @broadcast_splat_operand() -> tensor<5x5xi64> {
  %cst = arith.constant dense<5> : tensor<2xi64>
  %0 = "tf.BroadcastTo"(%cst, %cst) : (tensor<2xi64>, tensor<2xi64>) -> tensor<5x5xi64>
  func.return %0 : tensor<5x5xi64>
  // CHECK: %[[V0:.*]] = "tf.Const"() <{value = dense<5> : tensor<5x5xi64>}> : () -> tensor<5x5xi64>
  // CHECK: %[[V0]] : tensor<5x5xi64>

}

func.func @matrix_set_diag_v3_non_zero_k(%arg0: tensor<3x3xi64>, %arg1: tensor<3xi32>) -> tensor<3x3xi64> {
  %cst = arith.constant dense<1> : tensor<i32>
  %0 = "tf.MatrixSetDiagV3"(%arg0, %arg1, %cst) : (tensor<3x3xi64>, tensor<3xi32>, tensor<i32>) -> tensor<3x3xi64>
  func.return %0 : tensor<3x3xi64>

// CHECK-LABEL: @matrix_set_diag_v3_non_zero_k
// CHECK: tf.MatrixSetDiagV3
}

    -> (tensor<?x2xi64>, tensor<?x2xi64>, tensor<?x!tf_type.string>, tensor<?xi64>, tensor<2xi64>, tensor<2xi64>, tensor<32xf32>, tensor<32xf32>)

  // CHECK: tfrt.return {{.*}}, [[outputs]]#0
  func.return %outputs#0 : tensor<?x2xi64>
}

// CHECK-LABEL: func @no_native

    %1 = "tf.RngReadAndSkip"(%resource, %alg, %delta) : (tensor<!tf_type.resource<tensor<3xi64>>>, tensor<i32>, tensor<ui64>) -> tensor<3xi64>
    tf_device.return %1 : tensor<3xi64>
  }) : () -> tensor<3xi64>
  func.return %0 : tensor<3xi64>
}

// Test that ordering token resource of tf.CollectiveReduceV2 op is removed.
// CHECK-LABEL: func @decompose_collective_reduce_v2_op

        # CHECK: value = dense<{{\[\[}}1, 3, 2], [5, 4, 7]]> : tensor<2x3xi64>
      }
    }
  }
}
node {
  name: "z"
  op: "Const"
  device: "/device:CPU:0"
  attr {
    key: "dtype"
    value {
      type: DT_INT32
    }
  }
  attr {

  %7 = mhlo.constant dense<[0, 0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9, 9, 10, 10, 11, 11, 12, 12, 13, 13, 14, 14, 15]> : tensor<32xi64>
  %8 = "mhlo.broadcast_in_dim"(%7) <{broadcast_dimensions = dense<0> : tensor<1xi64>}> : (tensor<32xi64>) -> tensor<32x32xi64>
  %9 = mhlo.compare  LT, %8, %6 : (tensor<32x32xi64>, tensor<32x32xi64>) -> tensor<32x32xi1>
  %10 = mhlo.add %8, %4 : tensor<32x32xi64>

  } : (tensor<2x7xi64>) -> (tensor<2x7xi32>)
  %1 = "tf.Identity"(%0) : (tensor<2x7xi32>) -> tensor<2x7xi32>
  func.return %1 : tensor<2x7xi32>
}

// CHECK-LABEL: func @func
// CHECK-SAME: %arg0: tensor<2x7xi64> {mhlo.sharding = "\08\03\1A\03\01\01\02\22\02\00\01B\01\00"
func.func @func(%arg0: tensor<2x7xi64>) -> (tensor<2x7xi32>) {

  // CHECK-SAME: (tensor<3x3xi64>, tensor<3xi64>, tensor<i32>) -> tensor<3x3xi64>
}

// CHECK-LABEL: @testMatrixSetDiagV2
func.func @testMatrixSetDiagV2(%arg0: tensor<3x3xi64>, %arg1: tensor<3xi64>, %arg2: tensor<i32>) -> tensor<3x3xi64> {
  %0 = "tf.MatrixSetDiagV2"(%arg0, %arg1, %arg2) {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<3x3xi64>, tensor<3xi64>, tensor<i32>) -> tensor<3x3xi64>