commit 60832f6ee02a8ad4324c837ba5dd220606f6d52c
author Mihai Maruseac <mihaimaruseac@google.com> Fri Nov 19 09:04:55 2021 -0800
committer TensorFlower Gardener <gardener@tensorflow.org> Fri Nov 19 09:31:39 2021 -0800
tree 49706da0ab8f820257953884e652ae04bd781f54
parent e403fb82eded25c72d086ba992dcc81c6c98767c

Ensure `tree_complexity` is scalar in `BoostedTreesCalculateBestGainsPerFeature`

PiperOrigin-RevId: 411074185
Change-Id: Ida9f3c0ead019a6931a1157b0cfb161a9efa9e64

tensorflow/core/kernels/boosted_trees/stats_ops.cc

diff --git a/tensorflow/core/kernels/boosted_trees/stats_ops.cc b/tensorflow/core/kernels/boosted_trees/stats_ops.cc
index 503c192..1c85493 100644
--- a/tensorflow/core/kernels/boosted_trees/stats_ops.cc
+++ b/tensorflow/core/kernels/boosted_trees/stats_ops.cc
@@ -98,6 +98,10 @@
     const Tensor* tree_complexity_t;
     OP_REQUIRES_OK(context,
                    context->input("tree_complexity", &tree_complexity_t));
+    OP_REQUIRES(context, TensorShapeUtils::IsScalar(tree_complexity_t->shape()),
+                errors::InvalidArgument(
+                    "tree_complexity must be a scalar, got a tensor of shape ",
+                    tree_complexity_t->shape().DebugString()));
     const auto tree_complexity = tree_complexity_t->scalar<float>()();
     const Tensor* min_node_weight_t;
     OP_REQUIRES_OK(context,