shape_determination_fns.layout_preference_fn(
          cpu_tensor->shape(), cpu_tensor->dtype(), std::nullopt);
  TF_ASSIGN_OR_RETURN(xla::Shape shape,
                      shape_determination_fns.shape_representation_fn(
                          cpu_tensor->shape(), cpu_tensor->dtype(),
                          /*fast_mem=*/false, layout_preference));
  const xla::Layout* device_layout = &(shape.layout());

  for (auto operand_type_and_idx : llvm::enumerate(op.getOperandTypes())) {
    Type operand_type = operand_type_and_idx.value();
    int index = operand_type_and_idx.index();
    tensorflow::tpu::TPUCompileMetadataProto::Arg* arg = metadata->add_args();
    tensorflow::DataType dtype;
    tensorflow::Status status =
        tensorflow::ConvertToDataType(operand_type, &dtype);
    if (!status.ok())
      return op.emitOpError(

        absl::StrSplit(type, absl::ByAnyChar("()"));
    // If type has subtypes then parts[0] = type, parts[1] = subtypes,
    // parts[2] = ""
    if (parts.size() != 3 && parts.size() != 1) {
      return errors::InvalidArgument("Invalid type '", type, "'");
    } else if (parts.size() == 3) {
      // First part is the type, second is the subtype
      ArrayInfo::SubTypeInfo subtype;

    return mlir::TF::ShapeAttr::get(type.getContext(), tensor_type.getShape());
  }

  // If type is not a RankedTensor or UnrankedTensor, it must be a scalar.
  // Empty TensorShape indicates a scalar.
  return mlir::TF::ShapeAttr::get(type.getContext(), ArrayRef<int64_t>());
}

absl::StatusOr<TensorSpecProto> ConvertTypeToTensorSpecProto(
    const mlir::Type& type) {
  DataType dtype;

                                     const char* attr_name,
                                     const TF_DataType type,
                                     TF_Status* status) {
  tensorflow::DataType dtype = static_cast<tensorflow::DataType>(type);
  tensorflow::AddTypeConstraint(kernel_builder, attr_name, dtype, status);
}

void TF_KernelBuilder_HostMemory(TF_KernelBuilder* kernel_builder,

        xla_shape.dimensions(), &tensor_shape));
    arg_shapes.emplace_back(tensor_shape);

    DataType dtype;
    TF_RETURN_IF_ERROR(ConvertToDataType(input_type, &dtype));

    auto metadata_arg = metadata_proto.add_args();
    metadata_arg->set_kind(tpu::TPUCompileMetadataProto::Arg::PARAMETER);
    metadata_arg->set_dtype(dtype);
  }

  return absl::OkStatus();
}

absl::Status SetupReturnValues(mlir::ModuleOp module,

//     inputs: {
//       key  : "a"
//       value: {
//         name : "serving_default_a:0"
//         dtype: DT_FLOAT
//         tensor_shape: {
//         }
//       }
//     }
//     inputs: {
//       key  : "b"
//       value: {
//         name : "serving_default_b:0"
//         dtype: DT_FLOAT
//         tensor_shape: {
//         }
//       }
//     }
//     outputs: {

  }

  tensorflow::DataType dtype;
  auto status = tensorflow::ConvertToDataType(operand.getType(), &dtype);
  if (!status.ok()) {
    op->emitRemark() << "skipping legalization due to " << status.ToString();
    return tensorflow::XlaExpression::Invalid();
  }
  return tensorflow::XlaExpression::XlaOp(xla_op, dtype);
}

}  // namespace mhlo

    AllocatorAttributes attr;
    if (cc_input.dtype() == ::tensorflow::DT_VARIANT) {
      attr.set_on_host(true);
    }
    TF_RETURN_IF_ERROR(cc_ctx->allocate_temp(cc_input.dtype(), cc_input.shape(),
                                             cc_out, attr));

    switch (cc_input.dtype()) {
      case ::tensorflow::DT_INVALID: {
        *cc_out = Tensor(::tensorflow::DT_INVALID);

    }

    se::Stream* const src_compute_stream = src_xla_context->stream();
    TF_RET_CHECK(src_compute_stream != nullptr);
    TF_RET_CHECK(input->dtype() == output->dtype())
        << "input type: " << DataTypeString(input->dtype()) << " output type "
        << DataTypeString(output->dtype());
    TF_RET_CHECK(input->shape() == output->shape());
    TF_RET_CHECK(DMAHelper::CanUseDMA(input));
    auto* const src_compute_stream_impl =