Fix ubsan warning in TFE_Py_FastPathExecute_C

PiperOrigin-RevId: 234236587
This commit is contained in:
Akshay Modi 2019-02-15 16:41:32 -08:00 committed by TensorFlower Gardener
parent 27ca6327e7
commit 5dfe49e2d8

View File

@ -16,7 +16,6 @@ limitations under the License.
#include <cstring>
#include <thread>
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/python/eager/pywrap_tfe.h"
#include "absl/strings/str_cat.h"
@ -25,6 +24,7 @@ limitations under the License.
#include "tensorflow/c/c_api_internal.h"
#include "tensorflow/c/eager/c_api_internal.h"
#include "tensorflow/c/eager/tape.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/gtl/cleanup.h"
#include "tensorflow/core/lib/gtl/compactptrset.h"
#include "tensorflow/core/lib/gtl/flatmap.h"
@ -2448,14 +2448,14 @@ bool RaiseIfNotPySequence(PyObject* seq, const string& attr_name) {
bool RunCallbacks(
const FastPathOpExecInfo& op_exec_info, PyObject* args,
const std::vector<tensorflow::Safe_PyObjectPtr>& flattened_inputs,
const std::vector<tensorflow::Safe_PyObjectPtr>& flattened_attrs,
const std::vector<tensorflow::Safe_PyObjectPtr>* const flattened_inputs,
const std::vector<tensorflow::Safe_PyObjectPtr>* const flattened_attrs,
PyObject* flattened_result) {
if (!op_exec_info.run_callbacks) return true;
tensorflow::Safe_PyObjectPtr inputs(PyTuple_New(flattened_inputs.size()));
for (int i = 0; i < flattened_inputs.size(); i++) {
PyObject* input = flattened_inputs[i].get();
tensorflow::Safe_PyObjectPtr inputs(PyTuple_New(flattened_inputs->size()));
for (int i = 0; i < flattened_inputs->size(); i++) {
PyObject* input = (*flattened_inputs)[i].get();
Py_INCREF(input);
PyTuple_SET_ITEM(inputs.get(), i, input);
}
@ -2463,7 +2463,7 @@ bool RunCallbacks(
int num_non_inferred_attrs = PyTuple_GET_SIZE(args) -
op_exec_info.op_def->input_arg_size() -
kFastPathExecuteInputStartIndex;
int num_attrs = flattened_attrs.size() + num_non_inferred_attrs;
int num_attrs = flattened_attrs->size() + num_non_inferred_attrs;
tensorflow::Safe_PyObjectPtr attrs(PyTuple_New(num_attrs));
for (int i = 0; i < num_non_inferred_attrs; i++) {
@ -2475,7 +2475,7 @@ bool RunCallbacks(
}
for (int i = num_non_inferred_attrs; i < num_attrs; i++) {
PyObject* attr_or_name =
flattened_attrs.at(i - num_non_inferred_attrs).get();
flattened_attrs->at(i - num_non_inferred_attrs).get();
Py_INCREF(attr_or_name);
PyTuple_SET_ITEM(attrs.get(), i, attr_or_name);
}
@ -2795,8 +2795,8 @@ PyObject* TFE_Py_FastPathExecute_C(PyObject*, PyObject* args) {
PyList_SET_ITEM(flat_result.get(), i, EagerTensorFromHandle(retvals[i]));
}
if (!RunCallbacks(op_exec_info, args, *flattened_inputs, *flattened_attrs,
flat_result.get())) {
if (!RunCallbacks(op_exec_info, args, flattened_inputs.get(),
flattened_attrs.get(), flat_result.get())) {
return nullptr;
}