CherryPick:2.2 Fix tf.io.decode_raw bugs

This commit is contained in:
Mihai Maruseac 2021-04-30 06:36:59 -07:00 committed by Geeta Chavan
parent 5c401946a7
commit fe421f10e6

View File

@ -19,6 +19,7 @@ limitations under the License.
#include "tensorflow/core/framework/common_shape_fns.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/op_requires.h"
#include "tensorflow/core/framework/shape_inference.h"
namespace tensorflow {
@ -83,14 +84,13 @@ class DecodePaddedRawOp : public OpKernel {
// can copy the memory directly.
if (!convert_data_endianness_ || sizeof(T) == 1) {
for (int64 i = 0; i < flat_in.size(); ++i) {
const T* in_data = reinterpret_cast<const T*>(flat_in(i).data());
if (flat_in(i).size() > fixed_length) {
memcpy(out_data, in_data, fixed_length);
} else {
memcpy(out_data, in_data, flat_in(i).size());
}
out_data += fixed_length;
const auto to_copy =
std::min(flat_in(i).size(), static_cast<size_t>(fixed_length));
memcpy(out_data, flat_in(i).data(), to_copy);
// Note: increase out_data by width since it's already of type T* so
// each shift amount is implicitly multiplied by sizeof(T) according to
// pointer arithmetic rules.
out_data += width;
}
} else {
// Otherwise, the data is not in the host's byte order, and rather than a
@ -105,7 +105,10 @@ class DecodePaddedRawOp : public OpKernel {
p_in += sizeof(T), p_out += sizeof(T)) {
std::reverse_copy(p_in, p_in + sizeof(T), p_out);
}
out_data += fixed_length;
// Note: increase out_data by width since it's already of type T* so
// each shift amount is implicitly multiplied by sizeof(T) according to
// pointer arithmetic rules.
out_data += width;
}
}
}