mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-28 02:04:53 +08:00
631e2ee7a4
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/27136 make python udf serialization format to be binary plus tensor tables, so that tensors can be attached to autograd graph, handled in the same way as builtin operators ghstack-source-id: 91156141 Test Plan: unit tests Reviewed By: pritamdamania87 Differential Revision: D17405686 fbshipit-source-id: 4a8c9804f6ad239eb0655fa5daeb54580d4741fd
44 lines
1.2 KiB
C++
44 lines
1.2 KiB
C++
#pragma once
|
|
|
|
#include <torch/csrc/distributed/rpc/future_message.h>
|
|
#include <torch/csrc/distributed/rpc/message.h>
|
|
#include <torch/csrc/distributed/rpc/python_rpc_handler.h>
|
|
#include <torch/csrc/distributed/rpc/rpc_agent.h>
|
|
#include <torch/csrc/distributed/rpc/rref.h>
|
|
#include <torch/csrc/distributed/rpc/rref_context.h>
|
|
#include <torch/csrc/distributed/rpc/script_call.h>
|
|
#include <torch/csrc/distributed/rpc/script_remote_call.h>
|
|
#include <torch/csrc/distributed/rpc/script_ret.h>
|
|
#include <torch/csrc/jit/pybind_utils.h>
|
|
#include <torch/csrc/utils/pybind.h>
|
|
|
|
namespace torch {
|
|
namespace distributed {
|
|
namespace rpc {
|
|
|
|
py::object toPyObj(const Message& message);
|
|
|
|
std::shared_ptr<FutureMessage> pyRpcBuiltin(
|
|
RpcAgent& agent,
|
|
const WorkerId& dst,
|
|
const std::string& opName,
|
|
const py::args& args,
|
|
const py::kwargs& kwargs);
|
|
|
|
std::shared_ptr<FutureMessage> pyRpcPythonUdf(
|
|
RpcAgent& agent,
|
|
const WorkerId& dst,
|
|
const std::string& pickledPythonUDF,
|
|
std::vector<torch::Tensor>& tensors);
|
|
|
|
std::shared_ptr<RRef> pyRemoteBuiltin(
|
|
RpcAgent& agent,
|
|
const WorkerId& dst,
|
|
const std::string& opName,
|
|
const py::args& args,
|
|
const py::kwargs& kwargs);
|
|
|
|
} // namespace rpc
|
|
} // namespace distributed
|
|
} // namespace torch
|