Method Compile

tensorflow/compiler/tf2xla/kernels/async_io_ops.cc:50–63

}

void AsyncOutSendOp::Compile(XlaOpKernelContext* ctx) {
  // Always restore to the TensorShape.
  const Tensor& input = ctx->op_kernel_context()->input(0);
  const TensorShape& shape = input.shape();
  const DataType dtype = input.dtype();
  xla::Shape xla_shape;
  OP_REQUIRES_OK(ctx, TensorShapeToXLAShape(dtype, shape, &xla_shape));
  // Tensorflow always expects a row-major layout. So, always provide a shape
  // with row-major layout to AsyncOutSend. The XLA layout_assignment pass will
  // be constrained accordingly.
  xla::Shape shape_with_layout = xla::ShapeUtil::MakeShapeWithDescendingLayout(
      xla_shape.element_type(), xla_shape.dimensions());
  xla::AsyncOutSend(ctx->Input(0), shape_with_layout, key_);
}

REGISTER_XLA_OP(Name("_XlaAsyncOutSend"), AsyncOutSendOp);