// Converts a shape from MLIR to a TensorFlow tensor shape proto.
void ConvertToTensorShapeProto(llvm::ArrayRef<int64_t> shape,
                               TensorShapeProto* output_shape);

// Converts an MLIR type to a TensorFlow tensor shape.
PartialTensorShape ConvertTypeToTensorShape(const mlir::Type& type);

// Converts an MLIR shaped type to a TensorFlow shape attribute.

 public:
  XlaHostRecvDeviceContext(
      se::Stream* stream, const se::DeviceMemoryBase& device_memory_base,
      const xla::Shape& shape,
      tsl::AsyncValueRef<std::unique_ptr<se::Event>>& done_event)
      : stream_(stream),
        device_memory_base_(device_memory_base),
        shape_(shape),
        done_event_(done_event) {}

  void CopyCPUTensorToDevice(const Tensor* cpu_tensor, Device* device,

#include "xla/client/local_client.h"

namespace tensorflow {
namespace {
std::vector<const xla::Shape*> GetShapePointers(
    absl::Span<const xla::Shape> shapes) {
  std::vector<const xla::Shape*> shape_ptrs;
  shape_ptrs.reserve(shapes.size());
  for (const auto& shape : shapes) {
    shape_ptrs.push_back(&shape);
  }
  return shape_ptrs;
}
}  // namespace

absl::StatusOr<std::unique_ptr<xla::LocalExecutable>>

# RUN: tf-mlir-translate -graphdef-to-mlir -tf-enable-shape-inference-on-import=false %s -tf-input-arrays=p,x -tf-input-data-types="DT_INT32,DT_RESOURCE(DT_INT32)" -tf-output-arrays=p,x -o - | FileCheck %s -check-prefix=CHECK-NO-SHAPE

// RUN: tf-opt %s -tf-shape-inference=input-arg-shapes=1 -verify-diagnostics -split-input-file | FileCheck %s
// RUN: not tf-opt %s -tf-shape-inference=input-arg-shapes=* 2>&1 | FileCheck --check-prefix=INPUT_ARG_SHAPES_ERROR %s
// INPUT_ARG_SHAPES_ERROR: Missing input argument shapes

module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
    // CHECK-LABEL: func.func @main

    value {
      type: DT_INT32
    }
  }
  attr {
    key: "shape"
    value {
      shape {
        unknown_rank: true
      }
    }
  }
}
node {
  name: "x"
  op: "Placeholder"
  attr {
    key: "dtype"
    value {
      type: DT_VARIANT
    }
  }
  attr {
    key: "shape"
    value {
      shape {
        unknown_rank: true
      }
    }
  }
}
node {

// 7. Add necessary attributes to `node_def`, so we can replace it with a
//    XlaHostCompute node later. If all input shapes for XlaSendFromHost are
//    known, "shapes" attr will be set to the list of input shapes; otherwise
//    "shape_inference_graph" attr will be set to shape inference function name.
class RewriteOutsideCompilationSubgraphFn {
 public:
  RewriteOutsideCompilationSubgraphFn(

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

# RUN: not tf-mlir-translate -graphdef-to-mlir -tf-enable-shape-inference-on-import=false %s -tf-input-arrays=input -tf-input-data-types=DT_FLOAT -tf-input-shapes='' -tf-output-arrays=NotANodeInTheGraph -o - 2>&1 | FileCheck %s

# CHECK: Output NotANodeInTheGraph was not found in graph

node {
  name: "input"
  op: "Placeholder"
  device: "/device:CPU:0"
  attr {
    key: "dtype"
    value {
      type: DT_FLOAT
    }
  }
  attr {

// Check that attributes that define derived shapes are exported.

// CHECK: op: "PlaceholderWithDefault"
// CHECK: shape
// CHECK: unknown_rank: true
// CHECK: name: "static"
// CHECK: op: "PlaceholderWithDefault"
// CHECK: shape {
// CHECK-NEXT: }
// CHECK: name: "static_10"
// CHECK: op: "PlaceholderWithDefault"
// CHECK: shape
// CHECK: dim
// CHECK: size: 10

func.func @main() {