Function ReadVariableInputTensor

tensorflow/compiler/tf2xla/xla_op_kernel.cc:396–429

namespace {

Status ReadVariableInputTensor(const Tensor& tensor, DataType type,
                               const XlaOpKernelContext* ctx,
                               TensorShape* shape, xla::XlaOp* value) {
  const XlaExpression* expression = CastExpressionFromTensor(tensor);
  XlaResource* variable = expression->resource();
  TF_RET_CHECK(variable != nullptr);
  TF_RET_CHECK(variable->kind() == XlaResource::kVariable);
  if (!variable->initialized()) {
    return errors::FailedPrecondition("Read of uninitialized variable ",
                                      variable->name());
  }
  if (variable->type() != type) {
    return errors::InvalidArgument(
        "Type mismatch for read of variable ", variable->name(), ". Expected ",
        DataTypeString(type), "; got ", DataTypeString(variable->type()));
  }
  if (shape) {
    *shape = variable->shape();
  }

  TF_ASSIGN_OR_RETURN(xla::Shape representation_shape,
                      ctx->compiler()->options().shape_representation_fn(
                          variable->shape(), variable->type(),
                          /*use_fast_memory=*/false));
  xla::Shape xla_shape;
  TF_RETURN_IF_ERROR(
      TensorShapeToXLAShape(variable->type(), variable->shape(), &xla_shape));
  if (xla::ShapeUtil::Compatible(xla_shape, representation_shape)) {
    *value = variable->value();
  } else {
    *value = xla::Reshape(variable->value(), variable->shape().dim_sizes());
  }
  return Status::OK();
}

}  // namespace