Function AssignVariableTensor

tensorflow/compiler/tf2xla/xla_op_kernel.cc:536–566

namespace {

Status AssignVariableTensor(const Tensor& tensor, DataType type,
                            const XlaOpKernelContext* ctx, xla::XlaOp handle,
                            xla::XlaBuilder* builder) {
  const XlaExpression* expression = CastExpressionFromTensor(tensor);
  XlaResource* variable = expression->resource();
  TF_RET_CHECK(variable != nullptr);
  TF_RET_CHECK(variable->kind() == XlaResource::kVariable);

  auto shape_or_status = builder->GetShape(handle);
  if (!shape_or_status.ok()) {
    return shape_or_status.status();
  }
  TensorShape shape;
  TF_RETURN_IF_ERROR(
      XLAShapeToTensorShape(shape_or_status.ValueOrDie(), &shape));

  TF_RETURN_IF_ERROR(variable->SetTypeAndShape(type, shape));

  TF_ASSIGN_OR_RETURN(xla::Shape representation_shape,
                      ctx->compiler()->options().shape_representation_fn(
                          shape, type,
                          /*use_fast_memory=*/false));
  xla::Shape xla_shape;
  TF_RETURN_IF_ERROR(TensorShapeToXLAShape(type, shape, &xla_shape));
  if (!xla::ShapeUtil::Compatible(xla_shape, representation_shape)) {
    handle = xla::Reshape(handle,
                          xla::AsInt64Slice(representation_shape.dimensions()));
  }
  variable->SetRepresentationShape(representation_shape);
  return variable->SetValue(handle);
}

}  // namespace