// Copyright 2012 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// constant declarations

package const0

import "unsafe"

// constants declarations must be initialized by constants
var x = 0
const c0 = x /* ERROR "not constant" */

// typed constants must have constant types
const _ interface /* ERROR "invalid constant type" */ {} = 0

  %1 = stablehlo.transpose %0, dims = [0] : (tensor<2xf32>) -> tensor<2xf32>
  return %1 : tensor<2xf32>
}
// CHECK-DAG: %[[CONST_0:.+]] = stablehlo.constant dense<[0.000000e+00, 1.000000e+00]> : tensor<2xf32>
// CHECK-NOT: transpose
// CHECK: return %[[CONST_0]] : tensor<2xf32>

// -----

// CHECK-LABEL: transpose_simple_2d
func.func @transpose_simple_2d() -> tensor<3x2xf32> {

  // CHECK: "tf.EnqueueTPUEmbeddingSparseTensorBatch"(%[[ARG_0]], %[[ARG_1]], %[[ARG_2]], %[[ARG_3]], %[[ARG_4]], %[[ARG_5]], %[[CONST_0]], %[[CONST_0]], %[[CONST_0]], %[[CONST_MODE]])

    } : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<*xf32>
    return %0 : tensor<*xf32>
  }

// PerTensor-LABEL: func @depthwise_conv
// PerTensor-DAG: %[[CONST_0:.*]] = arith.constant dense<0.000000e+00> : tensor<6xf32>
// PerTensor: %[[CONST_1:.*]] = arith.constant dense
// PerTensor-NOT: tensor<2x3x1x6xf32>
// PerTensor-SAME: tensor<2x3x3x2xf32>

TEST(PerformStaticShapeInferenceBeforeEncapsulationTest, Basic) {
  // Build the graph:
  // "add" = "const_0" + "const_1"
  // "identity" = "add"
  tensorflow::Scope s = tensorflow::Scope::NewRootScope();
  Output const_0 = ops::Const(s.WithOpName("const_0"), 1, {2});
  Output const_1 = ops::Const(s.WithOpName("const_1"), 2, {2});
  Output add = ops::Add(s.WithOpName("add"), const_0, const_1);
  Output identity = ops::Identity(s.WithOpName("identity"), add);

  %2 = "tf.AddV2"(%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>

    return %0 : tensor<*xf32>
  }

// CHECK-LABEL: func @conv2d
// CHECK-DAG: %[[CONST_0:.*]] = arith.constant dense<0.000000e+00> : tensor<2xf32>
// CHECK-DAG: %[[CONST_1:.*]] = arith.constant dense<3.000000e+00> : tensor<2x3x512x2xf32>

func.func @main() -> tensor<2x3xi32> {
  %graph = tf_executor.graph {
    %0:2 = tf_executor.island wraps "tf.Const"() {value = dense<5> : tensor<2x3xi32>} : () -> tensor<2x3xi32> loc("Const0")
    %1:2 = tf_executor.island wraps "tf.Const"() {value = dense<4.2> : tensor<4x5xf32>} : () -> tensor<4x5xf32> loc("Const1")

    } : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<*xf32>
    return %0 : tensor<*xf32>
  }

// CHECK-LABEL: func @depthwise_conv
// CHECK-DAG: %[[CONST_0:.*]] = arith.constant dense<0.000000e+00> : tensor<6xf32>
// CHECK: %[[CONST_1:.*]] = arith.constant dense
// CHECK-NOT: tensor<2x3x3x2xf32>
// CHECK-SAME: tensor<2x3x1x6xf32>

// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() <{{{.*value = dense<"var_0"> : tensor<1x!tf_type.string>.*}}}>
// CHECK-DAG: %[[CONST_1:.*]] = "tf.Const"() <{{{.*value = dense<""> : tensor<1x!tf_type.string>.*}}}>
// CHECK: "tf.SaveV2"(%[[ARG]], %[[CONST_0]], %[[CONST_1]], %[[READ_VARIABLE]])
// CHECK: return
}

// -----