| 106 | } |
| 107 | |
| 108 | class GenEagerPythonOp : public python_op_gen_internal::GenPythonOp { |
| 109 | public: |
| 110 | GenEagerPythonOp(const OpDef& op_def, const ApiDef& api_def, |
| 111 | const string& function_name) |
| 112 | : python_op_gen_internal::GenPythonOp(op_def, api_def, function_name) { |
| 113 | op_name_ = function_name_; |
| 114 | absl::ConsumePrefix(&op_name_, "_"); |
| 115 | } |
| 116 | ~GenEagerPythonOp() override {} |
| 117 | |
| 118 | string Code() override; |
| 119 | |
| 120 | protected: |
| 121 | void HandleGraphMode(const string& function_setup); |
| 122 | |
| 123 | string GetEagerNotAllowedError(); |
| 124 | void ExpectListArg(const string& indentation, const string& arg_name, |
| 125 | string* output); |
| 126 | bool GetEagerFunctionSetup(const string& indentation, string* function_setup); |
| 127 | void GetOutputSizesAndNumOutputsExpr(std::vector<string>* output_sizes, |
| 128 | string* num_outputs_expr); |
| 129 | |
| 130 | void AddEagerFunctionTeardown(const string& indentation, |
| 131 | const std::vector<string>& output_sizes, |
| 132 | bool execute_record_gradient); |
| 133 | |
| 134 | bool AddEagerFastPathAndGraphCode(const string& parameters, |
| 135 | const std::vector<string>& output_sizes, |
| 136 | const string& eager_not_allowed_error); |
| 137 | bool AddEagerFallbackCode(const string& parameters, |
| 138 | const std::vector<string>& output_sizes, |
| 139 | const string& num_outputs_expr, |
| 140 | const string& eager_not_allowed_error); |
| 141 | void AddEagerFastPathExecute(); |
| 142 | |
| 143 | void AddEagerInferredAttrs(const string& indentation); |
| 144 | void AddEagerInputCasts(const string& indentation); |
| 145 | void AddEagerAttrs(const string& indentation); |
| 146 | void AddEagerExecute(const string& indentation, |
| 147 | const string& num_outputs_expr); |
| 148 | void AddDispatch(const string& prefix); |
| 149 | |
| 150 | void AddRawOpExport(const string& parameters); |
| 151 | |
| 152 | void AddAttrForArg(const string& attr, int arg_index) { |
| 153 | gtl::InsertIfNotPresent(&inferred_attrs_, attr, |
| 154 | op_def_.input_arg(arg_index).name()); |
| 155 | auto iter = attr_to_args_.find(attr); |
| 156 | if (iter == attr_to_args_.end()) { |
| 157 | attr_to_args_.insert(AttrToArgMap::value_type(attr, {arg_index})); |
| 158 | } else { |
| 159 | iter->second.push_back(arg_index); |
| 160 | } |
| 161 | } |
| 162 | |
| 163 | // Returns a string expression representing a flattened list of all |
| 164 | // the inputs given by `*input_indices` (or all inputs if |
| 165 | // `input_indices` is nullptr). `*output_sizes` can be used to unflatten. |