message: Prevent crash due to integer overflow followed by allocating negative sized array.
diff: @@ -13,6 +13,8 @@ See the License for the specific language governing permissions and
  limitations under the License.
  ==============================================================================*/
  +#include <limits>
  +
  #include "absl/container/flat_hash_map.h"
  #include "tensorflow/core/framework/op_kernel.h"
  #include "tensorflow/core/framework/op_requires.h"
  @@ -23,6 +25,9 @@ limitations under the License.
  namespace tensorflow {
  +// Don't allocate too large `BatchedMap<T>` objects
  +static int kMaxBatches = std::numeric_limits<int>::max();
  +
  template <class T>
  using BatchedMap = std::vector<absl::flat_hash_map<int64, T>>;
  @@ -236,6 +241,10 @@ class SparseCount : public OpKernel {
  bool is_1d = shape.NumElements() == 1;
  auto shape_vector = shape.flat<int64>();
  int num_batches = is_1d ? 1 : shape_vector(0);
  + OP_REQUIRES(
  + context, 0 < num_batches && num_batches < kMaxBatches,
  + errors::InvalidArgument("Cannot allocate ", num_batches,
  + " batches, is the dense shape too wide?"));
  for (int b = 0; b < shape_vector.size(); b++) {
  OP_REQUIRES(context, shape_vector(b) >= 0,
  