TensorSpec::TensorSpec()
    : shape_(std::initializer_list<int64_t>()), dtype_(DT_FLOAT) {}

TensorSpec::TensorSpec(PartialTensorShape shape, DataType dtype)
    : shape_(std::move(shape)), dtype_(dtype) {}

TensorSpec::TensorSpec(const TensorSpecProto& proto)
    : shape_(proto.shape()), dtype_(proto.dtype()) {}

const PartialTensorShape& TensorSpec::shape() const { return shape_; }

Keith Randall <******@****.***> 1627678827 -0700

      se::Stream* stream, 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) {}

  // Copies 'cpu_tensor' to `device_memory_base_` with `shape_`.
  // `device_tensor` is unused.

 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,

  func.return
}

// -----

func.func @testConcatOffest(%concat_dim: tensor<i32>, %shape0: tensor<3xi32>, %shape1: tensor<3xi32>) {
  // expected-error @+1 {{'tf.ConcatOffset' op requires operand and result 1 to have compatible shapes}}

    const Tensor* device_tensor, StringPiece tensor_name, Device* device,
    Tensor* cpu_tensor, StatusCallback done) {
  DataType dtype = EncodePrimitiveTypeAsDataType(shape_.element_type()).value();
  TensorShape tensor_shape;
  Status status = XLAShapeToTensorShape(shape_, &tensor_shape);
  if (!status.ok()) {
    done(status);
    return;
  }

  *cpu_tensor = Tensor(dtype, tensor_shape);

}

// Make invariant the `ShapeOp`s or a `ReadVariableOp` that's the `ShapeOp`'s
// predecessor.
void MakeShapeOpInvariant(tf_device::ReplicateOp replicate_op, int num_replicas,
                          Block* replicate_block, TF::ShapeOp shape_op) {
  // Ignore ShapeOps that have virtual devices.
  if (AncestorUsesVirtualDevice(replicate_op.getDevices(), shape_op)) return;

  Value input = shape_op.getInput();

# RUN: tf-mlir-translate -graphdef-to-mlir -tf-enable-unconditionally-use-set-output-shapes-on-import -tf-enable-shape-inference-on-import=false -tf-graph-as-function %s -o -  | FileCheck %s
# RUN: tf-mlir-translate -graphdef-to-mlir -tf-enable-unconditionally-use-set-output-shapes-on-import -tf-enable-shape-inference-on-import=true -tf-graph-as-function %s -o - | FileCheck %s

# Verify importing with _output_shapes enabled works as expected.

node {
  name: "_Arg"

Variable::Variable(ImmediateExecutionContext* ctx, DataType dtype,
                   TensorShape shape, absl::optional<std::string> name,
                   ImmediateTensorHandlePtr handle)
    : TensorHandleConvertible(std::move(handle)),
      name_(name.has_value() ? *name : "Variable"),
      dtype_(dtype),
      shape_(shape),
      ctx_(ctx) {}

Variable::~Variable() {

// CHECK-DAG: %[[SHAPE_OF:.*]] = shape.shape_of %[[CONV]] : tensor<?x3x2x2xf32> -> tensor<4xindex>
// CHECK-DAG: %[[BROADCAST:.*]] = stablehlo.dynamic_broadcast_in_dim %[[CONST_0]], %[[SHAPE_OF]], dims = [3] : (tensor<2xf32>, tensor<4xindex>) -> tensor<?x3x2x2xf32>