add sub reid

zhangmeng

2020-01-10 c3765bd24fe73747688a0ec2a550f219c9acb384

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
#pragma once
 
#include <torch/csrc/python_headers.h>
 
#include <torch/csrc/DynamicTypes.h>
#include <torch/csrc/THP.h>
#include <torch/csrc/autograd/variable.h>
#include <ATen/core/interned_strings.h>
#include <ATen/core/ivalue.h>
#include <torch/csrc/jit/pybind_utils.h>
#include <torch/csrc/jit/tracer.h>
#include <torch/csrc/utils/pybind.h>
 
#include <pybind11/functional.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
 
namespace py = pybind11;
 
namespace pybind11 {
namespace detail {
 
template <>
struct type_caster<torch::jit::IValue> {
 public:
  PYBIND11_TYPE_CASTER(torch::jit::IValue, _("IValue"));
 
  bool load(handle src, bool) {
    try {
      value = torch::jit::toIValue(src);
      return true;
    } catch (std::exception& e) {
      return false;
    }
  }
 
  static handle cast(
      torch::jit::IValue src,
      return_value_policy /* policy */,
      handle /* parent */) {
    return torch::jit::toPyObject(std::move(src)).release();
  }
};
 
template <>
struct type_caster<torch::jit::Symbol> {
 public:
  PYBIND11_TYPE_CASTER(torch::jit::Symbol, _("Symbol"));
 
  bool load(handle src, bool) {
    // TODO: Is there a way to py::cast that doesn't raise an exception on
    // failure?  Can we catch pybind11::cast_error here instead?
    std::string src_str;
    try {
      src_str = py::cast<std::string>(src);
    } catch (std::exception& e) {
      return false;
    }
    value = torch::jit::Symbol::fromQualString(src_str);
    return true;
  }
 
  static handle cast(
      torch::jit::Symbol src,
      return_value_policy /* policy */,
      handle /* parent */) {
    return py::cast(std::string(src.toQualString()), return_value_policy::copy)
        .release();
  }
};
 
template <>
struct type_caster<torch::jit::AttributeKind> {
 public:
  PYBIND11_TYPE_CASTER(torch::jit::AttributeKind, _("AttributeKind"));
 
  bool load(handle src, bool) {
    return false;
  }
 
  static handle cast(
      torch::jit::AttributeKind src,
      return_value_policy /* policy */,
      handle /* parent */) {
    return py::cast(
               std::string(torch::jit::toString(src)),
               return_value_policy::copy)
        .release();
  }
};
 
// See https://github.com/pybind/pybind11/issues/637
using ListCasterBase = pybind11::detail::
    list_caster<std::vector<torch::jit::Node*>, torch::jit::Node*>;
template <>
struct type_caster<std::vector<torch::jit::Node*>> : ListCasterBase {
  static handle cast(
      const std::vector<torch::jit::Node*>& src,
      return_value_policy,
      handle parent) {
    return ListCasterBase::cast(src, return_value_policy::reference, parent);
  }
  static handle cast(
      const std::vector<torch::jit::Node*>* src,
      return_value_policy pol,
      handle parent) {
    return cast(*src, pol, parent);
  }
};
 
} // namespace detail
} // namespace pybind11
 
namespace torch {
namespace jit {
 
static inline py::tuple tuple_tail(const py::tuple& tup) {
  py::tuple r(tup.size() - 1);
  for (size_t i = 1; i < tup.size(); i++) {
    r[i - 1] = tup[i];
  }
  return r;
}
 
} // namespace jit
} // namespace torch