mirrors/tensorflow
0
0 Fork 0
mirror of https://github.com/tensorflow/tensorflow.git synced 2026-10-26 00:09:51 +08:00
Code Activity

Harden TFLite CPU BatchMatMul Prepare ordering and quantization checks.

Browse Source 
- Move `InitializeTemporaries` in `batch_matmul::Prepare` to after the existing `lhs` and `rhs` rank (`>= 2` and `<= 5`) checks so `InitializeTemporaries` does not index `dims->data[rank - 2]` when `rank < 2`.
- Verify per-channel `quantized_dimension` in `EvalHybrid` before indexing `filter->dims->data[quantized_dimension]`, and propagate `EvalQuantized` errors in `Eval`.

PiperOrigin-RevId: 1005707488
This commit is contained in:
Changming Sun CA authored and TensorFlower Gardener committed 2026-10-23 10:56:06 -07:00
1 parent 9e381fa364
commit f780058909
2 files changed
+43 -4
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tensorflow/lite/kernels/batch_matmul.cc
+6 -4
View File
@@ -294,7 +294,6 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
OpContext op_context(context, node);
TF_LITE_ENSURE_OK(context, InitializeTemporaries(context, node, &op_context));
bool adj_x = op_context.params->adj_x;
bool adj_y = op_context.params->adj_y;
@@ -356,6 +355,8 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
TF_LITE_ENSURE(context, NumDimensions(rhs_data) >= 2);
TF_LITE_ENSURE(context, NumDimensions(rhs_data) <= 5);
TF_LITE_ENSURE_OK(context, InitializeTemporaries(context, node, &op_context));
const int lhs_rank = NumDimensions(lhs_data);
const int rhs_rank = NumDimensions(rhs_data);
const int output_rank = std::max(lhs_rank, rhs_rank);
@@ -498,6 +499,8 @@ TfLiteStatus EvalHybrid(TfLiteContext* context, TfLiteNode* node, OpData* data,
const auto* affine_quantization =
reinterpret_cast<TfLiteAffineQuantization*>(
filter->quantization.params);
TF_LITE_ENSURE_EQ(context, affine_quantization->quantized_dimension,
filter->dims->size - 2);
TF_LITE_ENSURE_EQ(
context, affine_quantization->scale->size,
filter->dims->data[affine_quantization->quantized_dimension]);
@@ -814,9 +817,8 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
break;
case kTfLiteInt8:
case kTfLiteInt16:
EvalQuantized(context, node, op_data, lhs_shape, lhs_tensor, rhs_shape,
rhs_tensor, output);
break;
return EvalQuantized(context, node, op_data, lhs_shape, lhs_tensor,
rhs_shape, rhs_tensor, output);
default:
TF_LITE_KERNEL_LOG(context,
"Currently BatchMatMul doesn't support type: %s",
tensorflow/lite/kernels/batch_matmul_test.cc
+37
View File
@@ -1175,5 +1175,42 @@ TEST(QuantizedBatchMatMulOpTest, SimpleTestQuantizedInt16Int8WithScales) {
ElementsAreArray(ArrayFloatNear({23, 23, 23, 57, 57, 57})));
}
class UnallocatedBatchMatMulOpModel : public SingleOpModel {
public:
UnallocatedBatchMatMulOpModel(const TensorData& lhs, const TensorData& rhs,
const TensorData& output) {
const int lhs_id = AddInput(lhs);
const int rhs_id = AddInput(rhs);
AddOutput(output);
SetBuiltinOp(BuiltinOperator_BATCH_MATMUL,
BuiltinOptions_BatchMatMulOptions,
CreateBatchMatMulOptions(builder_, /*adj_x=*/false,
/*adj_y=*/false,
/*asymmetric_quantize_inputs=*/true)
.Union());
BuildInterpreter({GetShape(lhs_id), GetShape(rhs_id)},
/*num_threads=*/1, /*allow_fp32_relax_to_fp16=*/false,
/*apply_delegate=*/false, /*allocate_and_delegate=*/false);
}
using SingleOpModel::AllocateTensors;
};
TEST(BatchMatMulOpTest, RejectsRank1InputInPrepare) {
UnallocatedBatchMatMulOpModel rank1_model(
/*lhs=*/{TensorType_FLOAT32, {3}},
/*rhs=*/{TensorType_INT8, {3, 4}, 0, 0, 1.0f, 0},
/*output=*/{TensorType_FLOAT32});
EXPECT_EQ(rank1_model.AllocateTensors(), kTfLiteError);
}
TEST(BatchMatMulOpTest, RejectsRank6InputInPrepare) {
UnallocatedBatchMatMulOpModel rank6_model(
/*lhs=*/{TensorType_FLOAT32, {1, 1, 1, 1, 2, 3}},
/*rhs=*/{TensorType_INT8, {1, 1, 1, 1, 3, 4}, 0, 0, 1.0f, 0},
/*output=*/{TensorType_FLOAT32});
EXPECT_EQ(rank6_model.AllocateTensors(), kTfLiteError);
}
} // namespace
} // namespace tflite