File tree 1 file changed +9
-0
lines changed
1 file changed +9
-0
lines changed Original file line number Diff line number Diff line change @@ -13,6 +13,8 @@ See the License for the specific language governing permissions and
1313limitations under the License.
1414==============================================================================*/
1515
16+ #include < limits>
17+
1618#include " absl/container/flat_hash_map.h"
1719#include " tensorflow/core/framework/op_kernel.h"
1820#include " tensorflow/core/framework/op_requires.h"
@@ -23,6 +25,9 @@ limitations under the License.
2325
2426namespace tensorflow {
2527
28+ // Don't allocate too large `BatchedMap<T>` objects
29+ static int kMaxBatches = std::numeric_limits<int >::max();
30+
2631template <class T >
2732using BatchedMap = std::vector<absl::flat_hash_map<int64_t , T>>;
2833
@@ -235,6 +240,10 @@ class SparseCount : public OpKernel {
235240
236241 bool is_1d = shape.NumElements () == 1 ;
237242 int num_batches = is_1d ? 1 : shape_vector (0 );
243+ OP_REQUIRES (
244+ context, 0 < num_batches && num_batches < kMaxBatches ,
245+ errors::InvalidArgument (" Cannot allocate " , num_batches,
246+ " batches, is the dense shape too wide?" ));
238247
239248 const auto values_values = values.flat <T>();
240249 const auto weight_values = weights.flat <W>();
You can’t perform that action at this time.
0 commit comments