Merge pull request #46505 from ddavis-2015:FloorDiv-pr3
PiperOrigin-RevId: 353845537 Change-Id: I62b245f6cb0428f27366ac7550892d510a217dd7
This commit is contained in:
TensorFlower Gardener
commit
bdde742906
170
tensorflow/lite/micro/kernels/floor_div.cc
Normal file
170
tensorflow/lite/micro/kernels/floor_div.cc
Normal file
@ -0,0 +1,170 @@
|
||||
/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
==============================================================================*/
|
||||
#include <math.h>
|
||||
#include <stddef.h>
|
||||
#include <stdint.h>
|
||||
|
||||
#include <functional>
|
||||
|
||||
#include "tensorflow/lite/c/common.h"
|
||||
#include "tensorflow/lite/kernels/internal/reference/binary_function.h"
|
||||
#include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
|
||||
#include "tensorflow/lite/kernels/internal/tensor.h"
|
||||
#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
|
||||
#include "tensorflow/lite/kernels/kernel_util.h"
|
||||
|
||||
namespace tflite {
|
||||
namespace ops {
|
||||
namespace builtin {
|
||||
namespace floor_div {
|
||||
namespace {
|
||||
|
||||
// Input/output tensor index.
|
||||
constexpr int kInputTensor1 = 0;
|
||||
constexpr int kInputTensor2 = 1;
|
||||
constexpr int kOutputTensor = 0;
|
||||
|
||||
// Op data for floor_div op.
|
||||
struct OpData {
|
||||
bool requires_broadcast;
|
||||
};
|
||||
|
||||
void* Init(TfLiteContext* context, const char* buffer, size_t length) {
|
||||
auto* data = new OpData;
|
||||
data->requires_broadcast = false;
|
||||
return data;
|
||||
}
|
||||
|
||||
void Free(TfLiteContext* context, void* buffer) {
|
||||
delete reinterpret_cast<OpData*>(buffer);
|
||||
}
|
||||
|
||||
TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
|
||||
TF_LITE_ENSURE_EQ(context, NumInputs(node), 2);
|
||||
TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
|
||||
|
||||
// Reinterprete the opaque data provided by user.
|
||||
OpData* data = reinterpret_cast<OpData*>(node->user_data);
|
||||
|
||||
const TfLiteTensor* input1;
|
||||
TF_LITE_ENSURE_OK(context,
|
||||
GetInputSafe(context, node, kInputTensor1, &input1));
|
||||
const TfLiteTensor* input2;
|
||||
TF_LITE_ENSURE_OK(context,
|
||||
GetInputSafe(context, node, kInputTensor2, &input2));
|
||||
TfLiteTensor* output;
|
||||
TF_LITE_ENSURE_OK(context,
|
||||
GetOutputSafe(context, node, kOutputTensor, &output));
|
||||
|
||||
TF_LITE_ENSURE_TYPES_EQ(context, input1->type, input2->type);
|
||||
|
||||
const TfLiteType type = input1->type;
|
||||
switch (type) {
|
||||
case kTfLiteFloat32:
|
||||
case kTfLiteInt32:
|
||||
break;
|
||||
default:
|
||||
TF_LITE_KERNEL_LOG(context, "Type '%s' is not supported by floor_div.",
|
||||
TfLiteTypeGetName(type));
|
||||
return kTfLiteError;
|
||||
}
|
||||
output->type = type;
|
||||
|
||||
data->requires_broadcast = !HaveSameShapes(input1, input2);
|
||||
|
||||
TfLiteIntArray* output_size = nullptr;
|
||||
if (data->requires_broadcast) {
|
||||
TF_LITE_ENSURE_OK(context, CalculateShapeForBroadcast(
|
||||
context, input1, input2, &output_size));
|
||||
} else {
|
||||
output_size = TfLiteIntArrayCopy(input1->dims);
|
||||
}
|
||||
|
||||
return context->ResizeTensor(context, output, output_size);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
TfLiteStatus EvalImpl(TfLiteContext* context, bool requires_broadcast,
|
||||
const TfLiteTensor* input1, const TfLiteTensor* input2,
|
||||
TfLiteTensor* output) {
|
||||
const T* denominator_data = GetTensorData<T>(input2);
|
||||
|
||||
// Validate the denominator.
|
||||
for (int i = 0; i < NumElements(input2); ++i) {
|
||||
if (std::equal_to<T>()(denominator_data[i], 0)) {
|
||||
TF_LITE_KERNEL_LOG(context, "Division by 0");
|
||||
return kTfLiteError;
|
||||
}
|
||||
}
|
||||
if (requires_broadcast) {
|
||||
reference_ops::BroadcastBinaryFunction4DSlow<T, T, T>(
|
||||
GetTensorShape(input1), GetTensorData<T>(input1),
|
||||
GetTensorShape(input2), denominator_data, GetTensorShape(output),
|
||||
GetTensorData<T>(output), reference_ops::FloorDiv<T>);
|
||||
} else {
|
||||
reference_ops::BinaryFunction<T, T, T>(
|
||||
GetTensorShape(input1), GetTensorData<T>(input1),
|
||||
GetTensorShape(input2), GetTensorData<T>(input2),
|
||||
GetTensorShape(output), GetTensorData<T>(output),
|
||||
reference_ops::FloorDiv<T>);
|
||||
}
|
||||
|
||||
return kTfLiteOk;
|
||||
}
|
||||
|
||||
TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
|
||||
OpData* data = reinterpret_cast<OpData*>(node->user_data);
|
||||
|
||||
const TfLiteTensor* input1;
|
||||
TF_LITE_ENSURE_OK(context,
|
||||
GetInputSafe(context, node, kInputTensor1, &input1));
|
||||
const TfLiteTensor* input2;
|
||||
TF_LITE_ENSURE_OK(context,
|
||||
GetInputSafe(context, node, kInputTensor2, &input2));
|
||||
TfLiteTensor* output;
|
||||
TF_LITE_ENSURE_OK(context,
|
||||
GetOutputSafe(context, node, kOutputTensor, &output));
|
||||
|
||||
switch (input1->type) {
|
||||
case kTfLiteInt32: {
|
||||
return EvalImpl<int32_t>(context, data->requires_broadcast, input1,
|
||||
input2, output);
|
||||
}
|
||||
case kTfLiteFloat32: {
|
||||
return EvalImpl<float>(context, data->requires_broadcast, input1, input2,
|
||||
output);
|
||||
}
|
||||
default: {
|
||||
TF_LITE_KERNEL_LOG(context, "Type '%s' is not supported by floor_div.",
|
||||
TfLiteTypeGetName(input1->type));
|
||||
return kTfLiteError;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace
|
||||
} // namespace floor_div
|
||||
|
||||
TfLiteRegistration* Register_FLOOR_DIV() {
|
||||
// Init, Free, Prepare, Eval are satisfying the Interface required by
|
||||
// TfLiteRegistration.
|
||||
static TfLiteRegistration r = {floor_div::Init, floor_div::Free,
|
||||
floor_div::Prepare, floor_div::Eval};
|
||||
return &r;
|
||||
}
|
||||
|
||||
} // namespace builtin
|
||||
} // namespace ops
|
||||
} // namespace tflite
|
||||
117
tensorflow/lite/micro/kernels/floor_div_test.cc
Normal file
117
tensorflow/lite/micro/kernels/floor_div_test.cc
Normal file
@ -0,0 +1,117 @@
|
||||
/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
==============================================================================*/
|
||||
#include <stdint.h>
|
||||
|
||||
#include <vector>
|
||||
|
||||
#include "tensorflow/lite/kernels/test_util.h"
|
||||
#include "tensorflow/lite/schema/schema_generated.h"
|
||||
|
||||
namespace tflite {
|
||||
namespace {
|
||||
|
||||
using ::testing::ElementsAre;
|
||||
|
||||
template <typename T>
|
||||
class FloorDivModel : public SingleOpModel {
|
||||
public:
|
||||
FloorDivModel(const TensorData& input1, const TensorData& input2,
|
||||
const TensorData& output) {
|
||||
input1_ = AddInput(input1);
|
||||
input2_ = AddInput(input2);
|
||||
output_ = AddOutput(output);
|
||||
SetBuiltinOp(BuiltinOperator_FLOOR_DIV, BuiltinOptions_FloorDivOptions,
|
||||
CreateFloorDivOptions(builder_).Union());
|
||||
BuildInterpreter({GetShape(input1_), GetShape(input2_)});
|
||||
}
|
||||
|
||||
int input1() { return input1_; }
|
||||
int input2() { return input2_; }
|
||||
|
||||
std::vector<T> GetOutput() { return ExtractVector<T>(output_); }
|
||||
std::vector<int> GetOutputShape() { return GetTensorShape(output_); }
|
||||
|
||||
private:
|
||||
int input1_;
|
||||
int input2_;
|
||||
int output_;
|
||||
};
|
||||
|
||||
TEST(FloorDivModel, Simple) {
|
||||
FloorDivModel<int32_t> model({TensorType_INT32, {1, 2, 2, 1}},
|
||||
{TensorType_INT32, {1, 2, 2, 1}},
|
||||
{TensorType_INT32, {}});
|
||||
model.PopulateTensor<int32_t>(model.input1(), {10, 9, 11, 3});
|
||||
model.PopulateTensor<int32_t>(model.input2(), {2, 2, 3, 4});
|
||||
model.Invoke();
|
||||
EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 2, 2, 1));
|
||||
EXPECT_THAT(model.GetOutput(), ElementsAre(5, 4, 3, 0));
|
||||
}
|
||||
|
||||
TEST(FloorDivModel, NegativeValue) {
|
||||
FloorDivModel<int32_t> model({TensorType_INT32, {1, 2, 2, 1}},
|
||||
{TensorType_INT32, {1, 2, 2, 1}},
|
||||
{TensorType_INT32, {}});
|
||||
model.PopulateTensor<int32_t>(model.input1(), {10, -9, -11, 7});
|
||||
model.PopulateTensor<int32_t>(model.input2(), {2, 2, -3, -4});
|
||||
model.Invoke();
|
||||
EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 2, 2, 1));
|
||||
EXPECT_THAT(model.GetOutput(), ElementsAre(5, -5, 3, -2));
|
||||
}
|
||||
|
||||
TEST(FloorDivModel, BroadcastFloorDiv) {
|
||||
FloorDivModel<int32_t> model({TensorType_INT32, {1, 2, 2, 1}},
|
||||
{TensorType_INT32, {1}}, {TensorType_INT32, {}});
|
||||
model.PopulateTensor<int32_t>(model.input1(), {10, -9, -11, 7});
|
||||
model.PopulateTensor<int32_t>(model.input2(), {-3});
|
||||
model.Invoke();
|
||||
EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 2, 2, 1));
|
||||
EXPECT_THAT(model.GetOutput(), ElementsAre(-4, 3, 3, -3));
|
||||
}
|
||||
|
||||
TEST(FloorDivModel, SimpleFloat) {
|
||||
FloorDivModel<float> model({TensorType_FLOAT32, {1, 2, 2, 1}},
|
||||
{TensorType_FLOAT32, {1, 2, 2, 1}},
|
||||
{TensorType_FLOAT32, {}});
|
||||
model.PopulateTensor<float>(model.input1(), {10.05, 9.09, 11.9, 3.01});
|
||||
model.PopulateTensor<float>(model.input2(), {2.05, 2.03, 3.03, 4.03});
|
||||
model.Invoke();
|
||||
EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 2, 2, 1));
|
||||
EXPECT_THAT(model.GetOutput(), ElementsAre(4.0, 4.0, 3.0, 0.0));
|
||||
}
|
||||
|
||||
TEST(FloorDivModel, NegativeValueFloat) {
|
||||
FloorDivModel<float> model({TensorType_FLOAT32, {1, 2, 2, 1}},
|
||||
{TensorType_FLOAT32, {1, 2, 2, 1}},
|
||||
{TensorType_FLOAT32, {}});
|
||||
model.PopulateTensor<float>(model.input1(), {10.03, -9.9, -11.0, 7.0});
|
||||
model.PopulateTensor<float>(model.input2(), {2.0, 2.3, -3.0, -4.1});
|
||||
model.Invoke();
|
||||
EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 2, 2, 1));
|
||||
EXPECT_THAT(model.GetOutput(), ElementsAre(5.0, -5.0, 3.0, -2.0));
|
||||
}
|
||||
|
||||
TEST(FloorDivModel, BroadcastFloorDivFloat) {
|
||||
FloorDivModel<float> model({TensorType_FLOAT32, {1, 2, 2, 1}},
|
||||
{TensorType_FLOAT32, {1}},
|
||||
{TensorType_FLOAT32, {}});
|
||||
model.PopulateTensor<float>(model.input1(), {10.03, -9.9, -11.0, 7.0});
|
||||
model.PopulateTensor<float>(model.input2(), {-3.3});
|
||||
model.Invoke();
|
||||
EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 2, 2, 1));
|
||||
EXPECT_THAT(model.GetOutput(), ElementsAre(-4.0, 2.0, 3.0, -3.0));
|
||||
}
|
||||
} // namespace
|
||||
} // namespace tflite
|
||||
Loading…
x
Reference in New Issue
Block a user