// CHECK-SAME: key = ""
  // CHECK-SAME: recv_key = "host_compute_channel_recv"
  // CHECK-SAME: send_key = "host_compute_channel_send"
  // CHECK-SAME: shape_inference_graph = @host_func
  // CHECK-SAME: shapes = [#tf_type.shape<*>, #tf_type.shape<3x?>]
  // CHECK-SAME: tpu_core = 0 : i64
  // CHECK: func @host_func
  // CHECK: %[[RECV_OUTPUT:[0-9]*]]:2 = "tf._XlaRecvAtHost"
  // CHECK-SAME: key = "host_compute_channel_send"

# RUN: tf-mlir-translate -graphdef-to-mlir -tf-enable-shape-inference-on-import=false %s -tf-graph-as-function | tf-mlir-translate -mlir-tf-graph-to-hlo-text -tf-input-shapes='' -tf-input-data-types=DT_FLOAT -tf-xla-emit-return-tuple | FileCheck %s

node {
  name: "arg"
  op: "_Arg"
  attr {
    key: "T"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "index"
    value {
      i: 0
    }
  }
}
node {

// CHECK-NEXT:       type: DT_FLOAT
// CHECK-NEXT:       type: DT_INT16

// CHECK:      attr {
// CHECK-NEXT:   key: "shapes"
// CHECK-NEXT:   value {
// CHECK-NEXT:     list {
// CHECK-NEXT:       shape {
// CHECK-NEXT:         dim {
// CHECK-NEXT:           size: 3
// CHECK:            shape {
// CHECK-NEXT:         dim {
// CHECK-NEXT:           size: 4
// CHECK-NEXT:         }
// CHECK-NEXT:         dim {

class TestModule(tf.Module):

  # Check that we get shapes annotated on function arguments.
  #
  # Besides checking the shape on the function input argument, this test also
  # checks that the shape on the input argument is propagated to the return
  # value.
  # We eventually want to move the shape inference to a pass separate from

// in a `shape` shaped tensor. Assumes row-major order. `indices` and `shape`
// should have the same size.
// Example: Index (2, 3) of a (4, 5)-shaped tensor has the contiguous offset of
// 2 * 5 + 3 = 13.
int64_t GetContiguousOffset(const ArrayRef<int64_t> indices,
                            const ArrayRef<int64_t> shape) {
  int64_t contiguous_offset = 0;
  int64_t base_offset = 1;

// is converted directly. If it is a shaped type, the element type is converted.
Status ConvertToDataType(mlir::Type type, DataType* dtype);

// Converts an TensorFlow shape to the one used in MLIR.
void ConvertToMlirShape(const TensorShape& input_shape,
                        llvm::SmallVectorImpl<int64_t>* shape);

// Converts an TensorFlow shape proto to the one used in MLIR.

node {
  name: "Placeholder"
  op: "Placeholder"
  attr {
    key: "dtype"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "shape"
    value {
      shape {
        dim {
          size: 2
        }

                           absl::string_view shapes,
                           GraphImportConfig::InputArrays* inputs);

Status ParseInputArrayInfo(
    const std::vector<string>& node_names,
    const std::vector<string>& node_dtypes,
    const std::vector<std::optional<std::vector<int>>>& node_shapes,
    GraphImportConfig::InputArrays* inputs);

// Parses shapes from the given string into shapes_vector which is a structured

namespace tensorflow {

llvm::SmallVector<int64_t> ConvertTFShapeToMlir(llvm::ArrayRef<int64_t> shapes);

llvm::SmallVector<int64_t> ConvertMlirShapeToTF(llvm::ArrayRef<int64_t> shape);

static constexpr int64_t kTFDynamicSize = -1;
mlir::RankedTensorType GetTypeFromTFTensorShape(llvm::ArrayRef<int64_t> shape,
                                                mlir::Type elementType,

  // CHECK-DAG:  %[[CSTR_LHS_SHAPE:.+]] = shape.shape_of %arg0
  // CHECK-DAG:  %[[CSTR_RHS_SHAPE:.+]] = shape.shape_of %arg1
  // CHECK-NEXT: %[[WITNESS:.+]] = shape.cstr_broadcastable %[[CSTR_LHS_SHAPE]], %[[CSTR_RHS_SHAPE]]
  // CHECK-NEXT: shape.assuming %[[WITNESS:.+]]
  // CHECK-DAG:    %[[LHS_SHAPE:.+]] = shape.shape_of %arg0
  // CHECK-DAG:    %[[RHS_SHAPE:.+]] = shape.shape_of %arg1