Status ModularFrameDecoder::DecodeGroup(const Rect& rect, BitReader* reader,
                                        int minShift, int maxShift,
                                        const ModularStreamId& stream,
                                        bool zerofill,
                                        PassesDecoderState* dec_state,
                                        ImageBundle* output) {
  JXL_DASSERT(stream.kind == ModularStreamId::kModularDC ||
              stream.kind == ModularStreamId::kModularAC);
  const size_t xsize = rect.xsize();
  const size_t ysize = rect.ysize();
  Image gi(xsize, ysize, full_image.bitdepth, 0);
  // start at the first bigger-than-groupsize non-metachannel
  size_t c = full_image.nb_meta_channels;
  for (; c < full_image.channel.size(); c++) {
    Channel& fc = full_image.channel[c];
    if (fc.w > frame_dim.group_dim || fc.h > frame_dim.group_dim) break;
  }
  size_t beginc = c;
  for (; c < full_image.channel.size(); c++) {
    Channel& fc = full_image.channel[c];
    int shift = std::min(fc.hshift, fc.vshift);
    if (shift > maxShift) continue;
    if (shift < minShift) continue;
    Rect r(rect.x0() >> fc.hshift, rect.y0() >> fc.vshift,
           rect.xsize() >> fc.hshift, rect.ysize() >> fc.vshift, fc.w, fc.h);
    if (r.xsize() == 0 || r.ysize() == 0) continue;
    Channel gc(r.xsize(), r.ysize());
    gc.hshift = fc.hshift;
    gc.vshift = fc.vshift;
    gi.channel.emplace_back(std::move(gc));
  }
  if (zerofill) {
    int gic = 0;
    for (c = beginc; c < full_image.channel.size(); c++) {
      Channel& fc = full_image.channel[c];
      int shift = std::min(fc.hshift, fc.vshift);
      if (shift > maxShift) continue;
      if (shift < minShift) continue;
      Rect r(rect.x0() >> fc.hshift, rect.y0() >> fc.vshift,
             rect.xsize() >> fc.hshift, rect.ysize() >> fc.vshift, fc.w, fc.h);
      if (r.xsize() == 0 || r.ysize() == 0) continue;
      for (size_t y = 0; y < r.ysize(); ++y) {
        pixel_type* const JXL_RESTRICT row_out = r.Row(&fc.plane, y);
        memset(row_out, 0, r.xsize() * sizeof(*row_out));
      }
      gic++;
    }
    return true;
  }
  ModularOptions options;
  if (!ModularGenericDecompress(
          reader, gi, /*header=*/nullptr, stream.ID(frame_dim), &options,
          /*undo_transforms=*/-1, &tree, &code, &context_map))
    return JXL_FAILURE("Failed to decode modular group");
  if (minShift == 0) {
    // Undo global transforms that have been pushed to the group level
    for (auto t : global_transform) {
      JXL_RETURN_IF_ERROR(t.Inverse(gi, global_header.wp_header));
    }
    if (!use_full_image) {
      JXL_RETURN_IF_ERROR(ModularImageToDecodedRect(
          gi, dec_state, nullptr, output, rect.Crop(dec_state->decoded)));
      return true;
    }
  }
  int gic = 0;
  for (c = beginc; c < full_image.channel.size(); c++) {
    Channel& fc = full_image.channel[c];
    int shift = std::min(fc.hshift, fc.vshift);
    if (shift > maxShift) continue;
    if (shift < minShift) continue;
    Rect r(rect.x0() >> fc.hshift, rect.y0() >> fc.vshift,
           rect.xsize() >> fc.hshift, rect.ysize() >> fc.vshift, fc.w, fc.h);
    if (r.xsize() == 0 || r.ysize() == 0) continue;
    CopyImageTo(/*rect_from=*/Rect(0, 0, r.xsize(), r.ysize()),
                /*from=*/gi.channel[gic].plane,
                /*rect_to=*/r, /*to=*/&fc.plane);
    gic++;
  }
  return true;
}

fix use_full_image==false case (#365)

Some fixes to the case where the full modular image is skipped:

- don't assume that everything happens at the modular AC group level (minShift==0), there can also be upsampling causing channels to have nonzero shift even when there's no Squeeze
- for partial decodes (when zerofill is true), don't try to fill the full image when it's not used. Instead initialize the decoded image with zeroes and skip the decoding.


Status ModularFrameDecoder::DecodeGroup(const Rect& rect, BitReader* reader,
                                        int minShift, int maxShift,
                                        const ModularStreamId& stream,
                                        bool zerofill,
                                        PassesDecoderState* dec_state,
                                        ImageBundle* output) {
  JXL_DASSERT(stream.kind == ModularStreamId::kModularDC ||
              stream.kind == ModularStreamId::kModularAC);
  const size_t xsize = rect.xsize();
  const size_t ysize = rect.ysize();
  Image gi(xsize, ysize, full_image.bitdepth, 0);
  // start at the first bigger-than-groupsize non-metachannel
  size_t c = full_image.nb_meta_channels;
  for (; c < full_image.channel.size(); c++) {
    Channel& fc = full_image.channel[c];
    if (fc.w > frame_dim.group_dim || fc.h > frame_dim.group_dim) break;
  }
  size_t beginc = c;
  for (; c < full_image.channel.size(); c++) {
    Channel& fc = full_image.channel[c];
    int shift = std::min(fc.hshift, fc.vshift);
    if (shift > maxShift) continue;
    if (shift < minShift) continue;
    Rect r(rect.x0() >> fc.hshift, rect.y0() >> fc.vshift,
           rect.xsize() >> fc.hshift, rect.ysize() >> fc.vshift, fc.w, fc.h);
    if (r.xsize() == 0 || r.ysize() == 0) continue;
    if (zerofill && use_full_image) {
      for (size_t y = 0; y < r.ysize(); ++y) {
        pixel_type* const JXL_RESTRICT row_out = r.Row(&fc.plane, y);
        memset(row_out, 0, r.xsize() * sizeof(*row_out));
      }
    } else {
      Channel gc(r.xsize(), r.ysize());
      if (zerofill) ZeroFillImage(&gc.plane);
      gc.hshift = fc.hshift;
      gc.vshift = fc.vshift;
      gi.channel.emplace_back(std::move(gc));
    }
  }
  if (zerofill && use_full_image) return true;
  ModularOptions options;
  if (!zerofill) {
    if (!ModularGenericDecompress(
            reader, gi, /*header=*/nullptr, stream.ID(frame_dim), &options,
            /*undo_transforms=*/-1, &tree, &code, &context_map)) {
      return JXL_FAILURE("Failed to decode modular group");
    }
  }
  // Undo global transforms that have been pushed to the group level
  if (!use_full_image) {
    for (auto t : global_transform) {
      JXL_RETURN_IF_ERROR(t.Inverse(gi, global_header.wp_header));
    }
    JXL_RETURN_IF_ERROR(ModularImageToDecodedRect(
        gi, dec_state, nullptr, output, rect.Crop(dec_state->decoded)));
    return true;
  }
  int gic = 0;
  for (c = beginc; c < full_image.channel.size(); c++) {
    Channel& fc = full_image.channel[c];
    int shift = std::min(fc.hshift, fc.vshift);
    if (shift > maxShift) continue;
    if (shift < minShift) continue;
    Rect r(rect.x0() >> fc.hshift, rect.y0() >> fc.vshift,
           rect.xsize() >> fc.hshift, rect.ysize() >> fc.vshift, fc.w, fc.h);
    if (r.xsize() == 0 || r.ysize() == 0) continue;
    JXL_ASSERT(use_full_image);
    CopyImageTo(/*rect_from=*/Rect(0, 0, r.xsize(), r.ysize()),
                /*from=*/gi.channel[gic].plane,
                /*rect_to=*/r, /*to=*/&fc.plane);
    gic++;
  }
  return true;
}


  void Compute(OpKernelContext* context) override {
    const Tensor& input = context->input(0);
    const TensorShape& input_shape = input.shape();
    const int32 input_dims = input_shape.dims();

    const Tensor& segment_id = context->input(1);
    const TensorShape& segment_id_shape = segment_id.shape();
    const int32 segment_dims = segment_id_shape.dims();

    const Tensor& num_segments_tensor = context->input(2);
    auto num_segments = num_segments_tensor.scalar<NUM_SEGMENTS_TYPE>()();

    OP_REQUIRES(context, segment_dims != 0,
                errors::InvalidArgument("Segment_id cannot have rank 0"));

    OP_REQUIRES(
        context, segment_dims <= input_dims,
        errors::OutOfRange("Invalid segment_id rank ", segment_dims,
                           " for input with ", input_dims, " dimension(s)"));
    for (auto i = 0; i < segment_dims; i++) {
      OP_REQUIRES(
          context, segment_id_shape.dim_size(i) == input_shape.dim_size(i),
          errors::InvalidArgument(
              "Segment dimension is ", segment_id_shape.dim_size(i),
              " while input dimension is ", input_dims, " in rank ", i));
    }

    // Making output tensor.
    Tensor* output_tensor = nullptr;
    TensorShape output_shape =
        GetOutputShape(input_shape, segment_id_shape, num_segments);
    OP_REQUIRES_OK(context, context->allocate_output("output", output_shape,
                                                     &output_tensor));

    // Preparating flat tensors.
    auto output_flat = output_tensor->flat<tstring>();
    auto flat_segment_id = segment_id.flat<INDICES_TYPE>();
    auto flat_input = input.flat<tstring>();

    for (int i = 0; i < flat_segment_id.size(); i++) {
      OP_REQUIRES(
          context,
          ((flat_segment_id(i) < num_segments) && (flat_segment_id(i) >= 0)),
          errors::InvalidArgument(
              "segment_ids are not allowed to exceed num_segments or"
              " to have negative values."));
    }

    int64 big_stride;
    int64 small_stride;
    std::tie(big_stride, small_stride) =
        GetStrides<INDICES_TYPE>(input_shape, segment_id_shape);
    auto relative_offset_set =
        GetFlattenedRelativeOffsets<INDICES_TYPE>(small_stride, big_stride);
    for (auto start_offset = 0; start_offset < big_stride; start_offset++) {
      for (auto i = 0; i < relative_offset_set.size(); i++) {
        auto output_index = start_offset + flat_segment_id(i) * big_stride;
        auto offset = start_offset + relative_offset_set[i];
        if (output_flat(output_index).length() != 0)
          output_flat(output_index).append(separator_.c_str());
        output_flat(output_index).append(flat_input(offset));
      }
    }
  }

Fix overflow CHECK issue with `tf.raw_ops.UnsortedSegmentJoin`.

PiperOrigin-RevId: 370766155
Change-Id: I33e7c6626224e1060a8a4ab51ad5d861c6d4c63e


  void Compute(OpKernelContext* context) override {
    const Tensor& input = context->input(0);
    const TensorShape& input_shape = input.shape();
    const int32 input_dims = input_shape.dims();

    const Tensor& segment_id = context->input(1);
    const TensorShape& segment_id_shape = segment_id.shape();
    const int32 segment_dims = segment_id_shape.dims();

    const Tensor& num_segments_tensor = context->input(2);
    OP_REQUIRES(context, num_segments_tensor.NumElements() != 0,
                errors::InvalidArgument("Number of segments cannot be empty."));
    auto num_segments = num_segments_tensor.scalar<NUM_SEGMENTS_TYPE>()();

    OP_REQUIRES(context, segment_dims != 0,
                errors::InvalidArgument("Segment_id cannot have rank 0"));

    OP_REQUIRES(
        context, segment_dims <= input_dims,
        errors::OutOfRange("Invalid segment_id rank ", segment_dims,
                           " for input with ", input_dims, " dimension(s)"));
    for (auto i = 0; i < segment_dims; i++) {
      OP_REQUIRES(
          context, segment_id_shape.dim_size(i) == input_shape.dim_size(i),
          errors::InvalidArgument(
              "Segment dimension is ", segment_id_shape.dim_size(i),
              " while input dimension is ", input_dims, " in rank ", i));
    }

    // Making output tensor.
    Tensor* output_tensor = nullptr;
    TensorShape output_shape =
        GetOutputShape(input_shape, segment_id_shape, num_segments);
    OP_REQUIRES_OK(context, context->allocate_output("output", output_shape,
                                                     &output_tensor));

    // Preparating flat tensors.
    auto output_flat = output_tensor->flat<tstring>();
    auto flat_segment_id = segment_id.flat<INDICES_TYPE>();
    auto flat_input = input.flat<tstring>();

    for (int i = 0; i < flat_segment_id.size(); i++) {
      OP_REQUIRES(
          context,
          ((flat_segment_id(i) < num_segments) && (flat_segment_id(i) >= 0)),
          errors::InvalidArgument(
              "segment_ids are not allowed to exceed num_segments or"
              " to have negative values."));
    }

    int64 big_stride;
    int64 small_stride;
    std::tie(big_stride, small_stride) =
        GetStrides<INDICES_TYPE>(input_shape, segment_id_shape);
    auto relative_offset_set =
        GetFlattenedRelativeOffsets<INDICES_TYPE>(small_stride, big_stride);
    for (auto start_offset = 0; start_offset < big_stride; start_offset++) {
      for (auto i = 0; i < relative_offset_set.size(); i++) {
        auto output_index = start_offset + flat_segment_id(i) * big_stride;
        auto offset = start_offset + relative_offset_set[i];
        if (output_flat(output_index).length() != 0)
          output_flat(output_index).append(separator_.c_str());
        output_flat(output_index).append(flat_input(offset));
      }
    }
  }


char *ogs_nas_5gs_suci_from_mobile_identity(
        ogs_nas_5gs_mobile_identity_t *mobile_identity)
{
    ogs_nas_5gs_mobile_identity_suci_t *mobile_identity_suci = NULL;
    ogs_plmn_id_t plmn_id;
    char tmp[OGS_MAX_IMSI_BCD_LEN+1];
    char routing_indicator[5];
    char *suci = NULL;

    ogs_assert(mobile_identity);

    mobile_identity_suci =
        (ogs_nas_5gs_mobile_identity_suci_t *)mobile_identity->buffer;
    ogs_assert(mobile_identity_suci);

    ogs_expect_or_return_val(mobile_identity_suci->h.supi_format ==
            OGS_NAS_5GS_SUPI_FORMAT_IMSI, NULL);
    ogs_expect_or_return_val(mobile_identity_suci->protection_scheme_id ==
            OGS_NAS_5GS_NULL_SCHEME, NULL);

    suci = ogs_msprintf("suci-%d-", mobile_identity_suci->h.supi_format);
    ogs_expect_or_return_val(suci, NULL);

    ogs_nas_to_plmn_id(&plmn_id, &mobile_identity_suci->nas_plmn_id);
    if (ogs_plmn_id_mnc_len(&plmn_id) == 2) {
        suci = ogs_mstrcatf(suci, "%03d-%02d-",
                ogs_plmn_id_mcc(&plmn_id), ogs_plmn_id_mnc(&plmn_id));
        ogs_expect_or_return_val(suci, NULL);
    } else {
        suci = ogs_mstrcatf(suci, "%03d-%03d-",
                ogs_plmn_id_mcc(&plmn_id), ogs_plmn_id_mnc(&plmn_id));
        ogs_expect_or_return_val(suci, NULL);
    }

    memset(routing_indicator, 0, sizeof(routing_indicator));
    if (mobile_identity_suci->routing_indicator1 != 0xf) {
        routing_indicator[0] =
            mobile_identity_suci->routing_indicator1 + '0';
        if (mobile_identity_suci->routing_indicator2 != 0xf) {
            routing_indicator[1] =
                mobile_identity_suci->routing_indicator2 + '0';
            if (mobile_identity_suci->routing_indicator3 != 0xf) {
                routing_indicator[2] =
                    mobile_identity_suci->routing_indicator3 + '0';
                if (mobile_identity_suci->routing_indicator4 != 0xf)
                    routing_indicator[3] =
                        mobile_identity_suci->routing_indicator4 + '0';
            }
        }
    }

    ogs_expect_or_return_val(mobile_identity->length > 8, NULL);
    ogs_buffer_to_bcd(mobile_identity_suci->scheme_output,
            mobile_identity->length - 8, tmp);

    suci = ogs_mstrcatf(suci, "%s-%d-%d-%s",
            routing_indicator,
            mobile_identity_suci->protection_scheme_id,
            mobile_identity_suci->home_network_pki_value,
            tmp);
    ogs_expect(suci);

    return suci;
}

[AMF] fix the crash when long MSIN length (#1206)

When the UE is in initially registered period,
if the length of MSIN(Part of SUPI) exceeds the normal length,
AMF stack smashing will be caused


char *ogs_nas_5gs_suci_from_mobile_identity(
        ogs_nas_5gs_mobile_identity_t *mobile_identity)
{
    ogs_nas_5gs_mobile_identity_suci_t *mobile_identity_suci = NULL;
    ogs_plmn_id_t plmn_id;
    char tmp[OGS_NAS_MAX_SCHEME_OUTPUT_LEN*2+1];
    char routing_indicator[5];
    char *suci = NULL;
    int scheme_output_len = 0;

    ogs_assert(mobile_identity);

    mobile_identity_suci =
        (ogs_nas_5gs_mobile_identity_suci_t *)mobile_identity->buffer;
    ogs_assert(mobile_identity_suci);

    ogs_expect_or_return_val(mobile_identity_suci->h.supi_format ==
            OGS_NAS_5GS_SUPI_FORMAT_IMSI, NULL);
    ogs_expect_or_return_val(mobile_identity_suci->protection_scheme_id ==
            OGS_NAS_5GS_NULL_SCHEME, NULL);

    suci = ogs_msprintf("suci-%d-", mobile_identity_suci->h.supi_format);
    ogs_expect_or_return_val(suci, NULL);

    ogs_nas_to_plmn_id(&plmn_id, &mobile_identity_suci->nas_plmn_id);
    if (ogs_plmn_id_mnc_len(&plmn_id) == 2) {
        suci = ogs_mstrcatf(suci, "%03d-%02d-",
                ogs_plmn_id_mcc(&plmn_id), ogs_plmn_id_mnc(&plmn_id));
        ogs_expect_or_return_val(suci, NULL);
    } else {
        suci = ogs_mstrcatf(suci, "%03d-%03d-",
                ogs_plmn_id_mcc(&plmn_id), ogs_plmn_id_mnc(&plmn_id));
        ogs_expect_or_return_val(suci, NULL);
    }

    memset(routing_indicator, 0, sizeof(routing_indicator));
    if (mobile_identity_suci->routing_indicator1 != 0xf) {
        routing_indicator[0] =
            mobile_identity_suci->routing_indicator1 + '0';
        if (mobile_identity_suci->routing_indicator2 != 0xf) {
            routing_indicator[1] =
                mobile_identity_suci->routing_indicator2 + '0';
            if (mobile_identity_suci->routing_indicator3 != 0xf) {
                routing_indicator[2] =
                    mobile_identity_suci->routing_indicator3 + '0';
                if (mobile_identity_suci->routing_indicator4 != 0xf)
                    routing_indicator[3] =
                        mobile_identity_suci->routing_indicator4 + '0';
            }
        }
    }

    scheme_output_len = mobile_identity->length - 8;
    ogs_expect_or_return_val(scheme_output_len > 0, NULL);
    ogs_expect_or_return_val(
            scheme_output_len <= OGS_NAS_MAX_SCHEME_OUTPUT_LEN, NULL);
    ogs_buffer_to_bcd(mobile_identity_suci->scheme_output,
            scheme_output_len, tmp);

    suci = ogs_mstrcatf(suci, "%s-%d-%d-%s",
            routing_indicator,
            mobile_identity_suci->protection_scheme_id,
            mobile_identity_suci->home_network_pki_value,
            tmp);
    ogs_expect(suci);

    return suci;
}


static int string_scan_range(RList *list, RBinFile *bf, int min,
			      const ut64 from, const ut64 to, int type, int raw, RBinSection *section) {
	RBin *bin = bf->rbin;
	ut8 tmp[R_STRING_SCAN_BUFFER_SIZE];
	ut64 str_start, needle = from;
	int count = 0, i, rc, runes;
	int str_type = R_STRING_TYPE_DETECT;

	// if list is null it means its gonna dump
	r_return_val_if_fail (bf, -1);

	if (type == -1) {
		type = R_STRING_TYPE_DETECT;
	}
	if (from == to) {
		return 0;
	}
	if (from > to) {
		eprintf ("Invalid range to find strings 0x%"PFMT64x" .. 0x%"PFMT64x"\n", from, to);
		return -1;
	}
	st64 len = (st64)(to - from);
	if (len < 1 || len > ST32_MAX) {
		eprintf ("String scan range is invalid (%"PFMT64d" bytes)\n", len);
		return -1;
	}
	ut8 *buf = calloc (len, 1);
	if (!buf || !min) {
		free (buf);
		return -1;
	}
	st64 vdelta = 0, pdelta = 0;
	RBinSection *s = NULL;
	bool ascii_only = false;
	PJ *pj = NULL;
	if (bf->strmode == R_MODE_JSON && !list) {
		pj = pj_new ();
		if (pj) {
			pj_a (pj);
		}
	}
	r_buf_read_at (bf->buf, from, buf, len);
	char *charset = r_sys_getenv ("RABIN2_CHARSET");
	if (!R_STR_ISEMPTY (charset)) {
		RCharset *ch = r_charset_new ();
		if (r_charset_use (ch, charset)) {
			int outlen = len * 4;
			ut8 *out = calloc (len, 4);
			if (out) {
				int res = r_charset_encode_str (ch, out, outlen, buf, len);
				int i;
				// TODO unknown chars should be translated to null bytes
				for (i = 0; i < res; i++) {
					if (out[i] == '?') {
						out[i] = 0;
					}
				}
				len = res;
				free (buf);
				buf = out;
			} else {
				eprintf ("Cannot allocate\n");
			}
		} else {
			eprintf ("Invalid value for RABIN2_CHARSET.\n");
		}
		r_charset_free (ch);
	}
	free (charset);
	RConsIsBreaked is_breaked = (bin && bin->consb.is_breaked)? bin->consb.is_breaked: NULL;
	// may oobread
	while (needle < to) {
		if (is_breaked && is_breaked ()) {
			break;
		}
		// smol optimization
		if (needle + 4 < to) {
			ut32 n1 = r_read_le32 (buf + needle - from);
			if (!n1) {
				needle += 4;
				continue;
			}
		}
		rc = r_utf8_decode (buf + needle - from, to - needle, NULL);
		if (!rc) {
			needle++;
			continue;
		}
		bool addr_aligned = !(needle % 4);

		if (type == R_STRING_TYPE_DETECT) {
			char *w = (char *)buf + needle + rc - from;
			if (((to - needle) > 8 + rc)) {
				// TODO: support le and be
				bool is_wide32le = (needle + rc + 2 < to) && (!w[0] && !w[1] && !w[2] && w[3] && !w[4]);
				// reduce false positives
				if (is_wide32le) {
					if (!w[5] && !w[6] && w[7] && w[8]) {
						is_wide32le = false;
					}
				}
				if (!addr_aligned) {
					is_wide32le = false;
				}
				///is_wide32be &= (n1 < 0xff && n11 < 0xff); // false; // n11 < 0xff;
				if (is_wide32le  && addr_aligned) {
					str_type = R_STRING_TYPE_WIDE32; // asume big endian,is there little endian w32?
				} else {
					// bool is_wide = (n1 && n2 && n1 < 0xff && (!n2 || n2 < 0xff));
					bool is_wide = needle + rc + 4 < to && !w[0] && w[1] && !w[2] && w[3] && !w[4];
					str_type = is_wide? R_STRING_TYPE_WIDE: R_STRING_TYPE_ASCII;
				}
			} else {
				if (rc > 1) {
					str_type = R_STRING_TYPE_UTF8; // could be charset if set :?
				} else {
					str_type = R_STRING_TYPE_ASCII;
				}
			}
		} else if (type == R_STRING_TYPE_UTF8) {
			str_type = R_STRING_TYPE_ASCII; // initial assumption
		} else {
			str_type = type;
		}
		runes = 0;
		str_start = needle;

		/* Eat a whole C string */
		for (i = 0; i < sizeof (tmp) - 4 && needle < to; i += rc) {
			RRune r = {0};
			if (str_type == R_STRING_TYPE_WIDE32) {
				rc = r_utf32le_decode (buf + needle - from, to - needle, &r);
				if (rc) {
					rc = 4;
				}
			} else if (str_type == R_STRING_TYPE_WIDE) {
				rc = r_utf16le_decode (buf + needle - from, to - needle, &r);
				if (rc == 1) {
					rc = 2;
				}
			} else {
				rc = r_utf8_decode (buf + needle - from, to - needle, &r);
				if (rc > 1) {
					str_type = R_STRING_TYPE_UTF8;
				}
			}

			/* Invalid sequence detected */
			if (!rc || (ascii_only && r > 0x7f)) {
				needle++;
				break;
			}

			needle += rc;

			if (r_isprint (r) && r != '\\') {
				if (str_type == R_STRING_TYPE_WIDE32) {
					if (r == 0xff) {
						r = 0;
					}
				}
				rc = r_utf8_encode (tmp + i, r);
				runes++;
				/* Print the escape code */
			} else if (r && r < 0x100 && strchr ("\b\v\f\n\r\t\a\033\\", (char)r)) {
				if ((i + 32) < sizeof (tmp) && r < 93) {
					tmp[i + 0] = '\\';
					tmp[i + 1] = "       abtnvfr             e  "
					             "                              "
					             "                              "
					             "  \\"[r];
				} else {
					// string too long
					break;
				}
				rc = 2;
				runes++;
			} else {
				/* \0 marks the end of C-strings */
				break;
			}
		}

		tmp[i++] = '\0';

		if (runes < min && runes >= 2 && str_type == R_STRING_TYPE_ASCII && needle < to) {
			// back up past the \0 to the last char just in case it starts a wide string
			needle -= 2;
		}
		if (runes >= min) {
			// reduce false positives
			int j, num_blocks, *block_list;
			int *freq_list = NULL, expected_ascii, actual_ascii, num_chars;
			if (str_type == R_STRING_TYPE_ASCII) {
				for (j = 0; j < i; j++) {
					char ch = tmp[j];
					if (ch != '\n' && ch != '\r' && ch != '\t') {
						if (!IS_PRINTABLE (tmp[j])) {
							continue;
						}
					}
				}
			}
			switch (str_type) {
			case R_STRING_TYPE_UTF8:
			case R_STRING_TYPE_WIDE:
			case R_STRING_TYPE_WIDE32:
				num_blocks = 0;
				block_list = r_utf_block_list ((const ut8*)tmp, i - 1,
						str_type == R_STRING_TYPE_WIDE? &freq_list: NULL);
				if (block_list) {
					for (j = 0; block_list[j] != -1; j++) {
						num_blocks++;
					}
				}
				if (freq_list) {
					num_chars = 0;
					actual_ascii = 0;
					for (j = 0; freq_list[j] != -1; j++) {
						num_chars += freq_list[j];
						if (!block_list[j]) { // ASCII
							actual_ascii = freq_list[j];
						}
					}
					free (freq_list);
					expected_ascii = num_blocks ? num_chars / num_blocks : 0;
					if (actual_ascii > expected_ascii) {
						ascii_only = true;
						needle = str_start;
						free (block_list);
						continue;
					}
				}
				free (block_list);
				if (num_blocks > R_STRING_MAX_UNI_BLOCKS) {
					needle++;
					continue;
				}
			}
			RBinString *bs = R_NEW0 (RBinString);
			if (!bs) {
				break;
			}
			bs->type = str_type;
			bs->length = runes;
			bs->size = needle - str_start;
			bs->ordinal = count++;
			// TODO: move into adjust_offset
			switch (str_type) {
			case R_STRING_TYPE_WIDE:
				if (str_start - from > 1) {
					const ut8 *p = buf + str_start - 2 - from;
					if (p[0] == 0xff && p[1] == 0xfe) {
						str_start -= 2; // \xff\xfe
					}
				}
				break;
			case R_STRING_TYPE_WIDE32:
				if (str_start - from > 3) {
					const ut8 *p = buf + str_start - 4 - from;
					if (p[0] == 0xff && p[1] == 0xfe) {
						str_start -= 4; // \xff\xfe\x00\x00
					}
				}
				break;
			}
			if (!s) {
				if (section) {
					s = section;
				} else if (bf->o) {
					s = r_bin_get_section_at (bf->o, str_start, false);
				}
				if (s) {
					vdelta = s->vaddr;
					pdelta = s->paddr;
				}
			}
			ut64 baddr = bf->loadaddr && bf->o? bf->o->baddr: bf->loadaddr;
			bs->paddr = str_start + baddr;
			bs->vaddr = str_start - pdelta + vdelta + baddr;
			bs->string = r_str_ndup ((const char *)tmp, i);
			if (list) {
				r_list_append (list, bs);
				if (bf->o) {
					ht_up_insert (bf->o->strings_db, bs->vaddr, bs);
				}
			} else {
				print_string (bf, bs, raw, pj);
				r_bin_string_free (bs);
			}
			if (from == 0 && to == bf->size) {
				/* force lookup section at the next one */
				s = NULL;
			}
		}
		ascii_only = false;
	}
	free (buf);
	if (pj) {
		pj_end (pj);
		if (bin) {
			RIO *io = bin->iob.io;
			if (io) {
				io->cb_printf ("%s", pj_string (pj));
			}
		}
		pj_free (pj);
	}
	return count;
}

Fix integer overflow in string search causing oobread ##crash

* Reported by @greatergoodest via huntrdev
* BountyID: 8a3dc5cb-08b3-4807-82b2-77f08c137a04
* Reproducer bfileovf


static int string_scan_range(RList *list, RBinFile *bf, int min,
			      const ut64 from, const ut64 to, int type, int raw, RBinSection *section) {
	RBin *bin = bf->rbin;
	ut8 tmp[R_STRING_SCAN_BUFFER_SIZE];
	ut64 str_start, needle = from;
	int count = 0, i, rc, runes;
	int str_type = R_STRING_TYPE_DETECT;

	// if list is null it means its gonna dump
	r_return_val_if_fail (bf, -1);

	if (type == -1) {
		type = R_STRING_TYPE_DETECT;
	}
	if (from == to) {
		return 0;
	}
	if (from > to) {
		eprintf ("Invalid range to find strings 0x%"PFMT64x" .. 0x%"PFMT64x"\n", from, to);
		return -1;
	}
	st64 len = (st64)(to - from);
	if (len < 1 || len > ST32_MAX) {
		eprintf ("String scan range is invalid (%"PFMT64d" bytes)\n", len);
		return -1;
	}
	ut8 *buf = calloc (len, 1);
	if (!buf || !min) {
		free (buf);
		return -1;
	}
	st64 vdelta = 0, pdelta = 0;
	RBinSection *s = NULL;
	bool ascii_only = false;
	PJ *pj = NULL;
	if (bf->strmode == R_MODE_JSON && !list) {
		pj = pj_new ();
		if (pj) {
			pj_a (pj);
		}
	}
	r_buf_read_at (bf->buf, from, buf, len);
	char *charset = r_sys_getenv ("RABIN2_CHARSET");
	if (!R_STR_ISEMPTY (charset)) {
		RCharset *ch = r_charset_new ();
		if (r_charset_use (ch, charset)) {
			int outlen = len * 4;
			ut8 *out = calloc (len, 4);
			if (out) {
				int res = r_charset_encode_str (ch, out, outlen, buf, len);
				int i;
				// TODO unknown chars should be translated to null bytes
				for (i = 0; i < res; i++) {
					if (out[i] == '?') {
						out[i] = 0;
					}
				}
				len = res;
				free (buf);
				buf = out;
			} else {
				eprintf ("Cannot allocate\n");
			}
		} else {
			eprintf ("Invalid value for RABIN2_CHARSET.\n");
		}
		r_charset_free (ch);
	}
	free (charset);
	RConsIsBreaked is_breaked = (bin && bin->consb.is_breaked)? bin->consb.is_breaked: NULL;
	// may oobread
	while (needle < to && needle < UT64_MAX - 4) {
		if (is_breaked && is_breaked ()) {
			break;
		}
		// smol optimization
		if (needle < to - 4) {
			ut32 n1 = r_read_le32 (buf + (needle - from));
			if (!n1) {
				needle += 4;
				continue;
			}
		}
		rc = r_utf8_decode (buf + (needle - from), to - needle, NULL);
		if (!rc) {
			needle++;
			continue;
		}
		bool addr_aligned = !(needle % 4);

		if (type == R_STRING_TYPE_DETECT) {
			char *w = (char *)buf + (needle + rc - from);
			if (((to - needle) > 8 + rc)) {
				// TODO: support le and be
				bool is_wide32le = (needle + rc + 2 < to) && (!w[0] && !w[1] && !w[2] && w[3] && !w[4]);
				// reduce false positives
				if (is_wide32le) {
					if (!w[5] && !w[6] && w[7] && w[8]) {
						is_wide32le = false;
					}
				}
				if (!addr_aligned) {
					is_wide32le = false;
				}
				///is_wide32be &= (n1 < 0xff && n11 < 0xff); // false; // n11 < 0xff;
				if (is_wide32le  && addr_aligned) {
					str_type = R_STRING_TYPE_WIDE32; // asume big endian,is there little endian w32?
				} else {
					// bool is_wide = (n1 && n2 && n1 < 0xff && (!n2 || n2 < 0xff));
					bool is_wide = needle + rc + 4 < to && !w[0] && w[1] && !w[2] && w[3] && !w[4];
					str_type = is_wide? R_STRING_TYPE_WIDE: R_STRING_TYPE_ASCII;
				}
			} else {
				if (rc > 1) {
					str_type = R_STRING_TYPE_UTF8; // could be charset if set :?
				} else {
					str_type = R_STRING_TYPE_ASCII;
				}
			}
		} else if (type == R_STRING_TYPE_UTF8) {
			str_type = R_STRING_TYPE_ASCII; // initial assumption
		} else {
			str_type = type;
		}
		runes = 0;
		str_start = needle;

		/* Eat a whole C string */
		for (i = 0; i < sizeof (tmp) - 4 && needle < to; i += rc) {
			RRune r = {0};
			if (str_type == R_STRING_TYPE_WIDE32) {
				rc = r_utf32le_decode (buf + needle - from, to - needle, &r);
				if (rc) {
					rc = 4;
				}
			} else if (str_type == R_STRING_TYPE_WIDE) {
				rc = r_utf16le_decode (buf + needle - from, to - needle, &r);
				if (rc == 1) {
					rc = 2;
				}
			} else {
				rc = r_utf8_decode (buf + (needle - from), to - needle, &r);
				if (rc > 1) {
					str_type = R_STRING_TYPE_UTF8;
				}
			}

			/* Invalid sequence detected */
			if (!rc || (ascii_only && r > 0x7f)) {
				needle++;
				break;
			}

			needle += rc;

			if (r_isprint (r) && r != '\\') {
				if (str_type == R_STRING_TYPE_WIDE32) {
					if (r == 0xff) {
						r = 0;
					}
				}
				rc = r_utf8_encode (tmp + i, r);
				runes++;
				/* Print the escape code */
			} else if (r && r < 0x100 && strchr ("\b\v\f\n\r\t\a\033\\", (char)r)) {
				if ((i + 32) < sizeof (tmp) && r < 93) {
					tmp[i + 0] = '\\';
					tmp[i + 1] = "       abtnvfr             e  "
					             "                              "
					             "                              "
					             "  \\"[r];
				} else {
					// string too long
					break;
				}
				rc = 2;
				runes++;
			} else {
				/* \0 marks the end of C-strings */
				break;
			}
		}

		tmp[i++] = '\0';

		if (runes < min && runes >= 2 && str_type == R_STRING_TYPE_ASCII && needle < to) {
			// back up past the \0 to the last char just in case it starts a wide string
			needle -= 2;
		}
		if (runes >= min) {
			// reduce false positives
			int j, num_blocks, *block_list;
			int *freq_list = NULL, expected_ascii, actual_ascii, num_chars;
			if (str_type == R_STRING_TYPE_ASCII) {
				for (j = 0; j < i; j++) {
					char ch = tmp[j];
					if (ch != '\n' && ch != '\r' && ch != '\t') {
						if (!IS_PRINTABLE (tmp[j])) {
							continue;
						}
					}
				}
			}
			switch (str_type) {
			case R_STRING_TYPE_UTF8:
			case R_STRING_TYPE_WIDE:
			case R_STRING_TYPE_WIDE32:
				num_blocks = 0;
				block_list = r_utf_block_list ((const ut8*)tmp, i - 1,
						str_type == R_STRING_TYPE_WIDE? &freq_list: NULL);
				if (block_list) {
					for (j = 0; block_list[j] != -1; j++) {
						num_blocks++;
					}
				}
				if (freq_list) {
					num_chars = 0;
					actual_ascii = 0;
					for (j = 0; freq_list[j] != -1; j++) {
						num_chars += freq_list[j];
						if (!block_list[j]) { // ASCII
							actual_ascii = freq_list[j];
						}
					}
					free (freq_list);
					expected_ascii = num_blocks ? num_chars / num_blocks : 0;
					if (actual_ascii > expected_ascii) {
						ascii_only = true;
						needle = str_start;
						free (block_list);
						continue;
					}
				}
				free (block_list);
				if (num_blocks > R_STRING_MAX_UNI_BLOCKS) {
					needle++;
					continue;
				}
			}
			RBinString *bs = R_NEW0 (RBinString);
			if (!bs) {
				break;
			}
			bs->type = str_type;
			bs->length = runes;
			bs->size = needle - str_start;
			bs->ordinal = count++;
			// TODO: move into adjust_offset
			switch (str_type) {
			case R_STRING_TYPE_WIDE:
				if (str_start - from > 1) {
					const ut8 *p = buf + str_start - 2 - from;
					if (p[0] == 0xff && p[1] == 0xfe) {
						str_start -= 2; // \xff\xfe
					}
				}
				break;
			case R_STRING_TYPE_WIDE32:
				if (str_start - from > 3) {
					const ut8 *p = buf + str_start - 4 - from;
					if (p[0] == 0xff && p[1] == 0xfe) {
						str_start -= 4; // \xff\xfe\x00\x00
					}
				}
				break;
			}
			if (!s) {
				if (section) {
					s = section;
				} else if (bf->o) {
					s = r_bin_get_section_at (bf->o, str_start, false);
				}
				if (s) {
					vdelta = s->vaddr;
					pdelta = s->paddr;
				}
			}
			ut64 baddr = bf->loadaddr && bf->o? bf->o->baddr: bf->loadaddr;
			bs->paddr = str_start + baddr;
			bs->vaddr = str_start - pdelta + vdelta + baddr;
			bs->string = r_str_ndup ((const char *)tmp, i);
			if (list) {
				r_list_append (list, bs);
				if (bf->o) {
					ht_up_insert (bf->o->strings_db, bs->vaddr, bs);
				}
			} else {
				print_string (bf, bs, raw, pj);
				r_bin_string_free (bs);
			}
			if (from == 0 && to == bf->size) {
				/* force lookup section at the next one */
				s = NULL;
			}
		}
		ascii_only = false;
	}
	free (buf);
	if (pj) {
		pj_end (pj);
		if (bin) {
			RIO *io = bin->iob.io;
			if (io) {
				io->cb_printf ("%s", pj_string (pj));
			}
		}
		pj_free (pj);
	}
	return count;
}


  void Compute(OpKernelContext* context) override {
    // Checks what we're remapping and inverts the relevant remapping Tensors to
    // be maps with key = old ID, value = new ID.
    std::unordered_map<int64_t, int64_t> old_row_to_new_row_map;
    std::vector<bool> row_id_present;
    const Tensor* row_remapping_t;
    OP_REQUIRES_OK(context, context->input("row_remapping", &row_remapping_t));
    const auto row_remapping = row_remapping_t->vec<int64_t>();
    OP_REQUIRES(context, row_remapping.size() == num_rows_,
                errors::InvalidArgument(strings::StrCat(
                    "Size of row_remapping is ", row_remapping.size(),
                    " instead of being equal to num_rows=", num_rows_)));
    OP_REQUIRES_OK(context, RemapVectorToMap(row_remapping, &row_id_present,
                                             &old_row_to_new_row_map));

    // Calculates the min/max old row ID that we need to read, to save us from
    // reading some unnecessary slices of the old tensor.
    int64_t min_old_row = -1;
    int64_t max_old_row = -1;
    for (int i = 0; i < row_remapping.size(); ++i) {
      if (min_old_row < 0 ||
          (row_remapping(i) >= 0 && row_remapping(i) < min_old_row)) {
        min_old_row = row_remapping(i);
      }
      if (max_old_row < 0 ||
          (row_remapping(i) >= 0 && row_remapping(i) > max_old_row)) {
        max_old_row = row_remapping(i);
      }
    }

    // Processes the remapping for columns.
    std::unordered_map<int64_t, int64_t> old_col_to_new_col_map;
    std::vector<bool> col_id_present;
    const Tensor* col_remapping_t;
    OP_REQUIRES_OK(context, context->input("col_remapping", &col_remapping_t));
    const auto col_remapping = col_remapping_t->vec<int64_t>();
    // Note that we always "remap rows", even when the row vocabulary does
    // not change, because partitioning requires a mapping from partitioned
    // Variables to the full checkpoints we load.
    const bool remap_cols = col_remapping.size() > 0;
    if (remap_cols) {
      OP_REQUIRES(
          context, col_remapping.size() == num_cols_,
          errors::InvalidArgument(strings::StrCat(
              "Provided col_remapping, but its size is ", col_remapping.size(),
              " instead of being equal to num_cols=", num_cols_)));
      OP_REQUIRES_OK(context, RemapVectorToMap(col_remapping, &col_id_present,
                                               &old_col_to_new_col_map));
    } else {
      col_id_present.clear();
      col_id_present.resize(num_cols_, true);
    }

    // Processes the checkpoint source and the provided Tensor name.
    const Tensor* ckpt_path_t;
    OP_REQUIRES_OK(context, context->input("ckpt_path", &ckpt_path_t));
    OP_REQUIRES(
        context, ckpt_path_t->NumElements() == 1,
        errors::InvalidArgument("The `ckpt_path` tensor must have exactly one "
                                "element, got tensor of shape ",
                                ckpt_path_t->shape().DebugString()));
    const string& ckpt_path = ckpt_path_t->scalar<tstring>()();
    const Tensor* old_tensor_name_t;
    OP_REQUIRES_OK(context,
                   context->input("old_tensor_name", &old_tensor_name_t));
    const string& old_tensor_name = old_tensor_name_t->scalar<tstring>()();

    LOG(INFO) << "Processing checkpoint : " << ckpt_path;
    BundleReader reader(context->env(), ckpt_path);
    OP_REQUIRES_OK(context, reader.status());

    DataType tensor_type;
    TensorShape tensor_shape;
    OP_REQUIRES_OK(context, reader.LookupDtypeAndShape(
                                old_tensor_name, &tensor_type, &tensor_shape));
    OP_REQUIRES(context, tensor_type == DT_FLOAT,
                errors::InvalidArgument(strings::StrCat(
                    "Tensor ", old_tensor_name, " has invalid type ",
                    DataTypeString(tensor_type), " instead of expected type ",
                    DataTypeString(DT_FLOAT))));
    // This op is limited to loading Tensors of rank 2 (matrices).
    OP_REQUIRES(
        context, tensor_shape.dims() == 2,
        errors::InvalidArgument(strings::StrCat(
            "Tensor ", old_tensor_name, " has shape ",
            tensor_shape.DebugString(), " of invalid rank ",
            tensor_shape.dims(), " instead of expected shape of rank 2.")));

    if (!remap_cols) {
      // TODO(weiho): Consider relaxing this restriction to allow partial column
      // loading (even when no column remapping is specified) if there turns out
      // to be a use case for it.
      OP_REQUIRES(context, num_cols_ == tensor_shape.dim_size(1),
                  errors::InvalidArgument(strings::StrCat(
                      "Tensor ", old_tensor_name, " has shape ",
                      tensor_shape.DebugString(),
                      ", where the size of its 2nd dimension is ",
                      tensor_shape.dim_size(1),
                      " instead of being equal to num_cols=", num_cols_)));
    }

    // Uses TensorSlice to potentially load the old tensor in chunks in case
    // memory usage is a concern.
    std::vector<TensorSlice> tensor_slices;
    TensorSlice slice(tensor_shape.dims());
    if (min_old_row >= 0 && max_old_row >= 0) {
      int64_t row_start = min_old_row;
      // TODO(weiho): Given the list of old row IDs of interest (the keys of
      // old_row_to_new_row_map), we could also try something smarter to
      // find some minimal set of covering ranges for the list of old row IDs
      // such that the size of each range is less than max_rows_in_memory_.
      while (row_start <= max_old_row) {
        const int64_t slice_length =
            max_rows_in_memory_ <= 0
                // If max_rows_in_memory_ <= 0, we just load the entire chunk.
                ? max_old_row - row_start + 1
                : std::min(max_rows_in_memory_, max_old_row - row_start + 1);
        slice.set_start(0, row_start);
        slice.set_length(0, slice_length);
        tensor_slices.push_back(slice);
        row_start += slice_length;
      }
    }

    // Allocates the output matrix.
    Tensor* output_matrix_t = nullptr;
    OP_REQUIRES_OK(context,
                   context->allocate_output("output_matrix",
                                            TensorShape({num_rows_, num_cols_}),
                                            &output_matrix_t));
    auto output_matrix = output_matrix_t->matrix<float>();

    // Iterates through tensor slices and copies over values from the old tensor
    // to the output matrix.
    int64_t row_index = min_old_row;
    int64_t rows_copied = 0;
    Tensor loaded_tensor_t;
    for (const TensorSlice& tensor_slice : tensor_slices) {
      LOG(INFO) << "Loading slice " << tensor_slice.DebugString();
      TensorShape slice_shape;
      OP_REQUIRES_OK(context,
                     tensor_slice.SliceTensorShape(tensor_shape, &slice_shape));
      // Potentially re-allocates the tensor buffer since the last slice may
      // have fewer rows than the other slices.
      if (loaded_tensor_t.shape() != slice_shape) {
        loaded_tensor_t = Tensor(DT_FLOAT, slice_shape);
      }
      OP_REQUIRES_OK(context, reader.LookupSlice(old_tensor_name, tensor_slice,
                                                 &loaded_tensor_t));

      // Iterates through the old loaded tensor slice row-by-row.
      for (int row = 0; row < loaded_tensor_t.dim_size(0); ++row, ++row_index) {
        if (row_index % 500000 == min_old_row) {
          LOG(INFO) << "Processing old row " << row_index;
        }

        // If the old row ID is not found in old_row_to_new_row_map, continue
        // to the next row; otherwise, copy it to the output matrix.
        const int64_t* new_row_ptr =
            gtl::FindOrNull(old_row_to_new_row_map, row_index);
        if (new_row_ptr == nullptr) {
          continue;
        }
        ++rows_copied;
        const int64_t new_row = *new_row_ptr;

        // Copies over the row element-by-element, in case remapping is needed
        // along the column axis.
        const auto& loaded_tensor = loaded_tensor_t.matrix<float>();
        for (int old_col = 0; old_col < loaded_tensor_t.dim_size(1);
             ++old_col) {
          int64_t new_col = old_col;
          if (remap_cols) {
            const int64_t* new_col_ptr =
                gtl::FindOrNull(old_col_to_new_col_map, old_col);
            if (new_col_ptr == nullptr) {
              // Column remapping is specified, but this column is not found in
              // old_col_to_new_col_map, so we leave it uninitialized, to be
              // filled in with initializing_values later.
              continue;
            }
            new_col = *new_col_ptr;
          }

          OP_REQUIRES(context,
                      new_row < num_rows_ && new_col < num_cols_ &&
                          new_row >= 0 && new_col >= 0,
                      errors::Internal(strings::StrCat(
                          "new_row=", new_row, " and new_col=", new_col,
                          " should have been less than num_rows_=", num_rows_,
                          " and num_cols_=", num_cols_,
                          " and non-negative. This should never have happened "
                          "if the code were correct. Please file a bug.")));
          output_matrix(new_row, new_col) = loaded_tensor(row, old_col);
        }
      }
    }
    LOG(INFO) << "Copied " << rows_copied << " rows from old matrix (with "
              << tensor_shape.dim_size(0) << " rows) to new matrix (with "
              << num_rows_ << " rows).";

    // At this point, there are potentially whole rows/columns uninitialized
    // (corresponding to the indices where row_id_present/col_id_present are
    // false). We fill this in cell-by-cell using row_id_present and
    // col_id_present while dequeuing from the initializing_values vector.
    const Tensor* initializing_values_t;
    OP_REQUIRES_OK(
        context, context->input("initializing_values", &initializing_values_t));
    const auto initializing_values = initializing_values_t->flat<float>();
    int64_t initializing_values_index = 0;
    for (int i = 0; i < num_rows_; ++i) {
      for (int j = 0; j < num_cols_; ++j) {
        if (row_id_present[i] && col_id_present[j]) continue;
        OP_REQUIRES(
            context, initializing_values_index < initializing_values.size(),
            errors::InvalidArgument(
                "initializing_values contained ", initializing_values.size(),
                " elements, but more missing values remain."));
        output_matrix(i, j) = initializing_values(initializing_values_index);
        ++initializing_values_index;
      }
    }

    // Checks that we used all the given initializing values.
    OP_REQUIRES(
        context, initializing_values_index == initializing_values.size(),
        errors::InvalidArgument(
            "initializing_values contained ", initializing_values.size(),
            " elements, but only ", initializing_values_index,
            " elements were used to fill in missing values."));
  }

Fix tf.raw_ops.LoadAndRemapMatrix vulnerability with invalid `row_remapping`.

Check that `row_remapping` has the correct dims().

PiperOrigin-RevId: 445522800


  void Compute(OpKernelContext* context) override {
    // Checks what we're remapping and inverts the relevant remapping Tensors to
    // be maps with key = old ID, value = new ID.
    std::unordered_map<int64_t, int64_t> old_row_to_new_row_map;
    std::vector<bool> row_id_present;
    const Tensor* row_remapping_t;
    OP_REQUIRES_OK(context, context->input("row_remapping", &row_remapping_t));
    OP_REQUIRES(
        context, row_remapping_t->dims() == 1,
        errors::InvalidArgument("The `row_remapping` tensor must be 1-D, got "
                                "a tensor of shape ",
                                row_remapping_t->shape().DebugString()));
    const auto row_remapping = row_remapping_t->vec<int64_t>();
    OP_REQUIRES(context, row_remapping.size() == num_rows_,
                errors::InvalidArgument(strings::StrCat(
                    "Size of row_remapping is ", row_remapping.size(),
                    " instead of being equal to num_rows=", num_rows_)));
    OP_REQUIRES_OK(context, RemapVectorToMap(row_remapping, &row_id_present,
                                             &old_row_to_new_row_map));

    // Calculates the min/max old row ID that we need to read, to save us from
    // reading some unnecessary slices of the old tensor.
    int64_t min_old_row = -1;
    int64_t max_old_row = -1;
    for (int i = 0; i < row_remapping.size(); ++i) {
      if (min_old_row < 0 ||
          (row_remapping(i) >= 0 && row_remapping(i) < min_old_row)) {
        min_old_row = row_remapping(i);
      }
      if (max_old_row < 0 ||
          (row_remapping(i) >= 0 && row_remapping(i) > max_old_row)) {
        max_old_row = row_remapping(i);
      }
    }

    // Processes the remapping for columns.
    std::unordered_map<int64_t, int64_t> old_col_to_new_col_map;
    std::vector<bool> col_id_present;
    const Tensor* col_remapping_t;
    OP_REQUIRES_OK(context, context->input("col_remapping", &col_remapping_t));
    const auto col_remapping = col_remapping_t->vec<int64_t>();
    // Note that we always "remap rows", even when the row vocabulary does
    // not change, because partitioning requires a mapping from partitioned
    // Variables to the full checkpoints we load.
    const bool remap_cols = col_remapping.size() > 0;
    if (remap_cols) {
      OP_REQUIRES(
          context, col_remapping.size() == num_cols_,
          errors::InvalidArgument(strings::StrCat(
              "Provided col_remapping, but its size is ", col_remapping.size(),
              " instead of being equal to num_cols=", num_cols_)));
      OP_REQUIRES_OK(context, RemapVectorToMap(col_remapping, &col_id_present,
                                               &old_col_to_new_col_map));
    } else {
      col_id_present.clear();
      col_id_present.resize(num_cols_, true);
    }

    // Processes the checkpoint source and the provided Tensor name.
    const Tensor* ckpt_path_t;
    OP_REQUIRES_OK(context, context->input("ckpt_path", &ckpt_path_t));
    OP_REQUIRES(
        context, ckpt_path_t->NumElements() == 1,
        errors::InvalidArgument("The `ckpt_path` tensor must have exactly one "
                                "element, got tensor of shape ",
                                ckpt_path_t->shape().DebugString()));
    const string& ckpt_path = ckpt_path_t->scalar<tstring>()();
    const Tensor* old_tensor_name_t;
    OP_REQUIRES_OK(context,
                   context->input("old_tensor_name", &old_tensor_name_t));
    const string& old_tensor_name = old_tensor_name_t->scalar<tstring>()();

    LOG(INFO) << "Processing checkpoint : " << ckpt_path;
    BundleReader reader(context->env(), ckpt_path);
    OP_REQUIRES_OK(context, reader.status());

    DataType tensor_type;
    TensorShape tensor_shape;
    OP_REQUIRES_OK(context, reader.LookupDtypeAndShape(
                                old_tensor_name, &tensor_type, &tensor_shape));
    OP_REQUIRES(context, tensor_type == DT_FLOAT,
                errors::InvalidArgument(strings::StrCat(
                    "Tensor ", old_tensor_name, " has invalid type ",
                    DataTypeString(tensor_type), " instead of expected type ",
                    DataTypeString(DT_FLOAT))));
    // This op is limited to loading Tensors of rank 2 (matrices).
    OP_REQUIRES(
        context, tensor_shape.dims() == 2,
        errors::InvalidArgument(strings::StrCat(
            "Tensor ", old_tensor_name, " has shape ",
            tensor_shape.DebugString(), " of invalid rank ",
            tensor_shape.dims(), " instead of expected shape of rank 2.")));

    if (!remap_cols) {
      // TODO(weiho): Consider relaxing this restriction to allow partial column
      // loading (even when no column remapping is specified) if there turns out
      // to be a use case for it.
      OP_REQUIRES(context, num_cols_ == tensor_shape.dim_size(1),
                  errors::InvalidArgument(strings::StrCat(
                      "Tensor ", old_tensor_name, " has shape ",
                      tensor_shape.DebugString(),
                      ", where the size of its 2nd dimension is ",
                      tensor_shape.dim_size(1),
                      " instead of being equal to num_cols=", num_cols_)));
    }

    // Uses TensorSlice to potentially load the old tensor in chunks in case
    // memory usage is a concern.
    std::vector<TensorSlice> tensor_slices;
    TensorSlice slice(tensor_shape.dims());
    if (min_old_row >= 0 && max_old_row >= 0) {
      int64_t row_start = min_old_row;
      // TODO(weiho): Given the list of old row IDs of interest (the keys of
      // old_row_to_new_row_map), we could also try something smarter to
      // find some minimal set of covering ranges for the list of old row IDs
      // such that the size of each range is less than max_rows_in_memory_.
      while (row_start <= max_old_row) {
        const int64_t slice_length =
            max_rows_in_memory_ <= 0
                // If max_rows_in_memory_ <= 0, we just load the entire chunk.
                ? max_old_row - row_start + 1
                : std::min(max_rows_in_memory_, max_old_row - row_start + 1);
        slice.set_start(0, row_start);
        slice.set_length(0, slice_length);
        tensor_slices.push_back(slice);
        row_start += slice_length;
      }
    }

    // Allocates the output matrix.
    Tensor* output_matrix_t = nullptr;
    OP_REQUIRES_OK(context,
                   context->allocate_output("output_matrix",
                                            TensorShape({num_rows_, num_cols_}),
                                            &output_matrix_t));
    auto output_matrix = output_matrix_t->matrix<float>();

    // Iterates through tensor slices and copies over values from the old tensor
    // to the output matrix.
    int64_t row_index = min_old_row;
    int64_t rows_copied = 0;
    Tensor loaded_tensor_t;
    for (const TensorSlice& tensor_slice : tensor_slices) {
      LOG(INFO) << "Loading slice " << tensor_slice.DebugString();
      TensorShape slice_shape;
      OP_REQUIRES_OK(context,
                     tensor_slice.SliceTensorShape(tensor_shape, &slice_shape));
      // Potentially re-allocates the tensor buffer since the last slice may
      // have fewer rows than the other slices.
      if (loaded_tensor_t.shape() != slice_shape) {
        loaded_tensor_t = Tensor(DT_FLOAT, slice_shape);
      }
      OP_REQUIRES_OK(context, reader.LookupSlice(old_tensor_name, tensor_slice,
                                                 &loaded_tensor_t));

      // Iterates through the old loaded tensor slice row-by-row.
      for (int row = 0; row < loaded_tensor_t.dim_size(0); ++row, ++row_index) {
        if (row_index % 500000 == min_old_row) {
          LOG(INFO) << "Processing old row " << row_index;
        }

        // If the old row ID is not found in old_row_to_new_row_map, continue
        // to the next row; otherwise, copy it to the output matrix.
        const int64_t* new_row_ptr =
            gtl::FindOrNull(old_row_to_new_row_map, row_index);
        if (new_row_ptr == nullptr) {
          continue;
        }
        ++rows_copied;
        const int64_t new_row = *new_row_ptr;

        // Copies over the row element-by-element, in case remapping is needed
        // along the column axis.
        const auto& loaded_tensor = loaded_tensor_t.matrix<float>();
        for (int old_col = 0; old_col < loaded_tensor_t.dim_size(1);
             ++old_col) {
          int64_t new_col = old_col;
          if (remap_cols) {
            const int64_t* new_col_ptr =
                gtl::FindOrNull(old_col_to_new_col_map, old_col);
            if (new_col_ptr == nullptr) {
              // Column remapping is specified, but this column is not found in
              // old_col_to_new_col_map, so we leave it uninitialized, to be
              // filled in with initializing_values later.
              continue;
            }
            new_col = *new_col_ptr;
          }

          OP_REQUIRES(context,
                      new_row < num_rows_ && new_col < num_cols_ &&
                          new_row >= 0 && new_col >= 0,
                      errors::Internal(strings::StrCat(
                          "new_row=", new_row, " and new_col=", new_col,
                          " should have been less than num_rows_=", num_rows_,
                          " and num_cols_=", num_cols_,
                          " and non-negative. This should never have happened "
                          "if the code were correct. Please file a bug.")));
          output_matrix(new_row, new_col) = loaded_tensor(row, old_col);
        }
      }
    }
    LOG(INFO) << "Copied " << rows_copied << " rows from old matrix (with "
              << tensor_shape.dim_size(0) << " rows) to new matrix (with "
              << num_rows_ << " rows).";

    // At this point, there are potentially whole rows/columns uninitialized
    // (corresponding to the indices where row_id_present/col_id_present are
    // false). We fill this in cell-by-cell using row_id_present and
    // col_id_present while dequeuing from the initializing_values vector.
    const Tensor* initializing_values_t;
    OP_REQUIRES_OK(
        context, context->input("initializing_values", &initializing_values_t));
    const auto initializing_values = initializing_values_t->flat<float>();
    int64_t initializing_values_index = 0;
    for (int i = 0; i < num_rows_; ++i) {
      for (int j = 0; j < num_cols_; ++j) {
        if (row_id_present[i] && col_id_present[j]) continue;
        OP_REQUIRES(
            context, initializing_values_index < initializing_values.size(),
            errors::InvalidArgument(
                "initializing_values contained ", initializing_values.size(),
                " elements, but more missing values remain."));
        output_matrix(i, j) = initializing_values(initializing_values_index);
        ++initializing_values_index;
      }
    }

    // Checks that we used all the given initializing values.
    OP_REQUIRES(
        context, initializing_values_index == initializing_values.size(),
        errors::InvalidArgument(
            "initializing_values contained ", initializing_values.size(),
            " elements, but only ", initializing_values_index,
            " elements were used to fill in missing values."));
  }


static av_always_inline void add_yblock(SnowContext *s, int sliced, slice_buffer *sb, IDWTELEM *dst, uint8_t *dst8, const uint8_t *obmc, int src_x, int src_y, int b_w, int b_h, int w, int h, int dst_stride, int src_stride, int obmc_stride, int b_x, int b_y, int add, int offset_dst, int plane_index){
    const int b_width = s->b_width  << s->block_max_depth;
    const int b_height= s->b_height << s->block_max_depth;
    const int b_stride= b_width;
    BlockNode *lt= &s->block[b_x + b_y*b_stride];
    BlockNode *rt= lt+1;
    BlockNode *lb= lt+b_stride;
    BlockNode *rb= lb+1;
    uint8_t *block[4];
    int tmp_step= src_stride >= 7*MB_SIZE ? MB_SIZE : MB_SIZE*src_stride;
    uint8_t *tmp = s->scratchbuf;
    uint8_t *ptmp;
    int x,y;

    if(b_x<0){
        lt= rt;
        lb= rb;
    }else if(b_x + 1 >= b_width){
        rt= lt;
        rb= lb;
    }
    if(b_y<0){
        lt= lb;
        rt= rb;
    }else if(b_y + 1 >= b_height){
        lb= lt;
        rb= rt;
    }

    if(src_x<0){ //FIXME merge with prev & always round internal width up to *16
        obmc -= src_x;
        b_w += src_x;
        if(!sliced && !offset_dst)
            dst -= src_x;
        src_x=0;
    }else if(src_x + b_w > w){
        b_w = w - src_x;
    }
    if(src_y<0){
        obmc -= src_y*obmc_stride;
        b_h += src_y;
        if(!sliced && !offset_dst)
            dst -= src_y*dst_stride;
        src_y=0;
    }else if(src_y + b_h> h){
        b_h = h - src_y;
    }

    if(b_w<=0 || b_h<=0) return;

    av_assert2(src_stride > 2*MB_SIZE + 5);

    if(!sliced && offset_dst)
        dst += src_x + src_y*dst_stride;
    dst8+= src_x + src_y*src_stride;
//    src += src_x + src_y*src_stride;

    ptmp= tmp + 3*tmp_step;
    block[0]= ptmp;
    ptmp+=tmp_step;
    ff_snow_pred_block(s, block[0], tmp, src_stride, src_x, src_y, b_w, b_h, lt, plane_index, w, h);

    if(same_block(lt, rt)){
        block[1]= block[0];
    }else{
        block[1]= ptmp;
        ptmp+=tmp_step;
        ff_snow_pred_block(s, block[1], tmp, src_stride, src_x, src_y, b_w, b_h, rt, plane_index, w, h);
    }

    if(same_block(lt, lb)){
        block[2]= block[0];
    }else if(same_block(rt, lb)){
        block[2]= block[1];
    }else{
        block[2]= ptmp;
        ptmp+=tmp_step;
        ff_snow_pred_block(s, block[2], tmp, src_stride, src_x, src_y, b_w, b_h, lb, plane_index, w, h);
    }

    if(same_block(lt, rb) ){
        block[3]= block[0];
    }else if(same_block(rt, rb)){
        block[3]= block[1];
    }else if(same_block(lb, rb)){
        block[3]= block[2];
    }else{
        block[3]= ptmp;
        ff_snow_pred_block(s, block[3], tmp, src_stride, src_x, src_y, b_w, b_h, rb, plane_index, w, h);
    }
    if(sliced){
        s->dwt.inner_add_yblock(obmc, obmc_stride, block, b_w, b_h, src_x,src_y, src_stride, sb, add, dst8);
    }else{
        for(y=0; y<b_h; y++){
            //FIXME ugly misuse of obmc_stride
            const uint8_t *obmc1= obmc + y*obmc_stride;
            const uint8_t *obmc2= obmc1+ (obmc_stride>>1);
            const uint8_t *obmc3= obmc1+ obmc_stride*(obmc_stride>>1);
            const uint8_t *obmc4= obmc3+ (obmc_stride>>1);
            for(x=0; x<b_w; x++){
                int v=   obmc1[x] * block[3][x + y*src_stride]
                        +obmc2[x] * block[2][x + y*src_stride]
                        +obmc3[x] * block[1][x + y*src_stride]
                        +obmc4[x] * block[0][x + y*src_stride];

                v <<= 8 - LOG2_OBMC_MAX;
                if(FRAC_BITS != 8){
                    v >>= 8 - FRAC_BITS;
                }
                if(add){
                    v += dst[x + y*dst_stride];
                    v = (v + (1<<(FRAC_BITS-1))) >> FRAC_BITS;
                    if(v&(~255)) v= ~(v>>31);
                    dst8[x + y*src_stride] = v;
                }else{
                    dst[x + y*dst_stride] -= v;
                }
            }
        }
    }
}

avcodec/snow: split block clipping checks

Fixes out of array read
Fixes: d4476f68ca1c1c57afbc45806f581963-asan_heap-oob_2266b27_8607_cov_4044577381_snow_chroma_bug.avi
Found-by: Mateusz "j00ru" Jurczyk and Gynvael Coldwind
Signed-off-by: Michael Niedermayer <michaelni@gmx.at>


static av_always_inline void add_yblock(SnowContext *s, int sliced, slice_buffer *sb, IDWTELEM *dst, uint8_t *dst8, const uint8_t *obmc, int src_x, int src_y, int b_w, int b_h, int w, int h, int dst_stride, int src_stride, int obmc_stride, int b_x, int b_y, int add, int offset_dst, int plane_index){
    const int b_width = s->b_width  << s->block_max_depth;
    const int b_height= s->b_height << s->block_max_depth;
    const int b_stride= b_width;
    BlockNode *lt= &s->block[b_x + b_y*b_stride];
    BlockNode *rt= lt+1;
    BlockNode *lb= lt+b_stride;
    BlockNode *rb= lb+1;
    uint8_t *block[4];
    int tmp_step= src_stride >= 7*MB_SIZE ? MB_SIZE : MB_SIZE*src_stride;
    uint8_t *tmp = s->scratchbuf;
    uint8_t *ptmp;
    int x,y;

    if(b_x<0){
        lt= rt;
        lb= rb;
    }else if(b_x + 1 >= b_width){
        rt= lt;
        rb= lb;
    }
    if(b_y<0){
        lt= lb;
        rt= rb;
    }else if(b_y + 1 >= b_height){
        lb= lt;
        rb= rt;
    }

    if(src_x<0){ //FIXME merge with prev & always round internal width up to *16
        obmc -= src_x;
        b_w += src_x;
        if(!sliced && !offset_dst)
            dst -= src_x;
        src_x=0;
    }
    if(src_x + b_w > w){
        b_w = w - src_x;
    }
    if(src_y<0){
        obmc -= src_y*obmc_stride;
        b_h += src_y;
        if(!sliced && !offset_dst)
            dst -= src_y*dst_stride;
        src_y=0;
    }
    if(src_y + b_h> h){
        b_h = h - src_y;
    }

    if(b_w<=0 || b_h<=0) return;

    av_assert2(src_stride > 2*MB_SIZE + 5);

    if(!sliced && offset_dst)
        dst += src_x + src_y*dst_stride;
    dst8+= src_x + src_y*src_stride;
//    src += src_x + src_y*src_stride;

    ptmp= tmp + 3*tmp_step;
    block[0]= ptmp;
    ptmp+=tmp_step;
    ff_snow_pred_block(s, block[0], tmp, src_stride, src_x, src_y, b_w, b_h, lt, plane_index, w, h);

    if(same_block(lt, rt)){
        block[1]= block[0];
    }else{
        block[1]= ptmp;
        ptmp+=tmp_step;
        ff_snow_pred_block(s, block[1], tmp, src_stride, src_x, src_y, b_w, b_h, rt, plane_index, w, h);
    }

    if(same_block(lt, lb)){
        block[2]= block[0];
    }else if(same_block(rt, lb)){
        block[2]= block[1];
    }else{
        block[2]= ptmp;
        ptmp+=tmp_step;
        ff_snow_pred_block(s, block[2], tmp, src_stride, src_x, src_y, b_w, b_h, lb, plane_index, w, h);
    }

    if(same_block(lt, rb) ){
        block[3]= block[0];
    }else if(same_block(rt, rb)){
        block[3]= block[1];
    }else if(same_block(lb, rb)){
        block[3]= block[2];
    }else{
        block[3]= ptmp;
        ff_snow_pred_block(s, block[3], tmp, src_stride, src_x, src_y, b_w, b_h, rb, plane_index, w, h);
    }
    if(sliced){
        s->dwt.inner_add_yblock(obmc, obmc_stride, block, b_w, b_h, src_x,src_y, src_stride, sb, add, dst8);
    }else{
        for(y=0; y<b_h; y++){
            //FIXME ugly misuse of obmc_stride
            const uint8_t *obmc1= obmc + y*obmc_stride;
            const uint8_t *obmc2= obmc1+ (obmc_stride>>1);
            const uint8_t *obmc3= obmc1+ obmc_stride*(obmc_stride>>1);
            const uint8_t *obmc4= obmc3+ (obmc_stride>>1);
            for(x=0; x<b_w; x++){
                int v=   obmc1[x] * block[3][x + y*src_stride]
                        +obmc2[x] * block[2][x + y*src_stride]
                        +obmc3[x] * block[1][x + y*src_stride]
                        +obmc4[x] * block[0][x + y*src_stride];

                v <<= 8 - LOG2_OBMC_MAX;
                if(FRAC_BITS != 8){
                    v >>= 8 - FRAC_BITS;
                }
                if(add){
                    v += dst[x + y*dst_stride];
                    v = (v + (1<<(FRAC_BITS-1))) >> FRAC_BITS;
                    if(v&(~255)) v= ~(v>>31);
                    dst8[x + y*src_stride] = v;
                }else{
                    dst[x + y*dst_stride] -= v;
                }
            }
        }
    }
}


  void Compute(OpKernelContext *ctx) override {
    const Tensor *indices_t, *values_t, *shape_t, *reduction_axes_t;
    OP_REQUIRES_OK(ctx, ctx->input("input_indices", &indices_t));
    OP_REQUIRES_OK(ctx, ctx->input("input_values", &values_t));
    OP_REQUIRES_OK(ctx, ctx->input("input_shape", &shape_t));
    OP_REQUIRES_OK(ctx, ctx->input("reduction_axes", &reduction_axes_t));

    OP_REQUIRES_OK(ctx, ValidateInputs(shape_t, reduction_axes_t));

    // TODO(zongheng): we will call Reorder() below, which will modify
    // in-place the underlying indices and values buffers.  To avoid
    // surprises of this kernel being stateful, we work around the above by
    // making deep copies here.  Remove this if/when we change Reorder()'s
    // semantics.
    const auto shape_vec = shape_t->vec<int64>();
    SparseTensor sp;
    OP_REQUIRES_OK(ctx, SparseTensor::Create(
        tensor::DeepCopy(*indices_t), tensor::DeepCopy(*values_t),
                    TensorShape(shape_vec), &sp));
    ReduceDetails reduction = SparseTensorReduceHelper(
        sp, reduction_axes_t->flat<int32>(), keep_dims_);

    Tensor *out_values;
    OP_REQUIRES_OK(
        ctx, ctx->allocate_output(0, reduction.reduced_shape, &out_values));
    auto out_flat = out_values->flat<T>();
    out_flat.setZero();

    Tensor tmp_reduced_val;
    OP_REQUIRES_OK(ctx, ctx->allocate_temp(DataTypeToEnum<T>::value,
                                           TensorShape({}), &tmp_reduced_val));
    auto reduced_val = tmp_reduced_val.scalar<T>();

    // Compute strides, and use it to convert coords to flat index.  The
    // coordinates returned by .group() have the same ndims as group_by_dims.
    gtl::InlinedVector<int64, 8> output_strides(reduction.group_by_dims.size());
    if (!output_strides.empty()) {  // Do this iff we don't reduce all.
      output_strides.back() = 1;
      for (int d = output_strides.size() - 2; d >= 0; --d) {
        output_strides[d] =
            output_strides[d + 1] * shape_vec(reduction.group_by_dims[d + 1]);
      }
    }

    auto CoordinatesToFlatIndex = [](ArraySlice<int64> coords,
                                     ArraySlice<int64> strides) -> int64 {
      if (strides.empty()) {  // Reduce all.
        return 0;
      }
      CHECK_EQ(coords.size(), strides.size());
      int64_t idx = 0;
      for (int i = 0; i < coords.size(); ++i) {
        idx += coords[i] * strides[i];
      }
      return idx;
    };

    // Each group maps one-on-one onto a value in the reduced tensor.
    // g.group() provides the coordinates of a particular reduced value.
    sp.Reorder<T>(reduction.reorder_dims);
    for (const auto &g : sp.group(reduction.group_by_dims)) {
      Op::template Run<T>(ctx, reduced_val, g.template values<T>());
      const int64_t idx = CoordinatesToFlatIndex(g.group(), output_strides);
      out_flat(idx) = reduced_val();
      VLOG(2) << "coords: " << absl::StrJoin(g.group(), ",")
              << "; idx: " << idx << "; group " << Op::Name() << ": "
              << reduced_val();
    }
  }

Prevent heap OOB in sparse reduction ops.

PiperOrigin-RevId: 387934524
Change-Id: I894aa30f1e454f09b471d565b4a325da49322c1a


  void Compute(OpKernelContext *ctx) override {
    const Tensor *indices_t, *values_t, *shape_t, *reduction_axes_t;
    OP_REQUIRES_OK(ctx, ctx->input("input_indices", &indices_t));
    OP_REQUIRES_OK(ctx, ctx->input("input_values", &values_t));
    OP_REQUIRES_OK(ctx, ctx->input("input_shape", &shape_t));
    OP_REQUIRES_OK(ctx, ctx->input("reduction_axes", &reduction_axes_t));

    OP_REQUIRES_OK(ctx, ValidateInputs(shape_t, reduction_axes_t));

    // TODO(zongheng): we will call Reorder() below, which will modify
    // in-place the underlying indices and values buffers.  To avoid
    // surprises of this kernel being stateful, we work around the above by
    // making deep copies here.  Remove this if/when we change Reorder()'s
    // semantics.
    const auto shape_vec = shape_t->vec<int64>();
    SparseTensor sp;
    OP_REQUIRES_OK(ctx, SparseTensor::Create(
        tensor::DeepCopy(*indices_t), tensor::DeepCopy(*values_t),
                    TensorShape(shape_vec), &sp));
    ReduceDetails reduction = SparseTensorReduceHelper(
        sp, reduction_axes_t->flat<int32>(), keep_dims_);

    Tensor *out_values;
    OP_REQUIRES_OK(
        ctx, ctx->allocate_output(0, reduction.reduced_shape, &out_values));
    auto out_flat = out_values->flat<T>();
    out_flat.setZero();

    Tensor tmp_reduced_val;
    OP_REQUIRES_OK(ctx, ctx->allocate_temp(DataTypeToEnum<T>::value,
                                           TensorShape({}), &tmp_reduced_val));
    auto reduced_val = tmp_reduced_val.scalar<T>();

    // Compute strides, and use it to convert coords to flat index.  The
    // coordinates returned by .group() have the same ndims as group_by_dims.
    gtl::InlinedVector<int64, 8> output_strides(reduction.group_by_dims.size());
    if (!output_strides.empty()) {  // Do this iff we don't reduce all.
      output_strides.back() = 1;
      for (int d = output_strides.size() - 2; d >= 0; --d) {
        output_strides[d] =
            output_strides[d + 1] * shape_vec(reduction.group_by_dims[d + 1]);
      }
    }

    auto CoordinatesToFlatIndex = [](ArraySlice<int64> coords,
                                     ArraySlice<int64> strides) -> int64 {
      if (strides.empty()) {  // Reduce all.
        return 0;
      }
      CHECK_EQ(coords.size(), strides.size());
      int64_t idx = 0;
      for (int i = 0; i < coords.size(); ++i) {
        idx += coords[i] * strides[i];
      }
      return idx;
    };

    // Each group maps one-on-one onto a value in the reduced tensor.
    // g.group() provides the coordinates of a particular reduced value.
    sp.Reorder<T>(reduction.reorder_dims);
    for (const auto &g : sp.group(reduction.group_by_dims)) {
      Op::template Run<T>(ctx, reduced_val, g.template values<T>());
      OP_REQUIRES(ctx,
                  output_strides.empty() ||
                  (g.group().size() == output_strides.size()),
                  errors::Internal(
                      "Expected group size and output_strides size to match",
                      ", but got ", g.group().size(), " and ",
                      output_strides.size()));
      const int64_t idx = CoordinatesToFlatIndex(g.group(), output_strides);
      OP_REQUIRES(ctx,
                  idx >= 0 && idx < out_flat.size(),
                  errors::Internal(
                      "Obtained a write index of ", idx,
                      " which is outside of bounds of [0, ",
                      out_flat.size(), ")"));
      out_flat(idx) = reduced_val();
      VLOG(2) << "coords: " << absl::StrJoin(g.group(), ",")
              << "; idx: " << idx << "; group " << Op::Name() << ": "
              << reduced_val();
    }
  }


  void Compute(OpKernelContext* context) override {
    // Here's the basic idea:
    // Batch and depth dimension are independent from row and col dimension. And
    // because FractionalAvgPool currently only support pooling along row and
    // col, we can basically think of this 4D tensor backpropagation as
    // operation of a series of 2D planes.
    //
    // For each element of a 'slice' (2D plane) of output_backprop, we need to
    // figure out its contributors when doing FractionalAvgPool operation. This
    // can be done based on row_pooling_sequence, col_pooling_seq and
    // overlapping.
    // Once we figure out the original contributors, we just need to evenly
    // divide the value of this element among these contributors.
    //
    // Internally, we divide the out_backprop tensor and store it in a temporary
    // tensor of double type. And cast it to the corresponding type.
    typedef Eigen::Map<const Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>>
        ConstEigenMatrixMap;
    typedef Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>>
        EigenDoubleMatrixMap;

    // Grab the inputs.
    const Tensor& orig_input_tensor_shape = context->input(0);
    OP_REQUIRES(context,
                orig_input_tensor_shape.dims() == 1 &&
                    orig_input_tensor_shape.NumElements() == 4,
                errors::InvalidArgument("original input tensor shape must be"
                                        "1-dimensional and 4 elements"));
    const Tensor& out_backprop = context->input(1);
    const Tensor& row_seq_tensor = context->input(2);
    const Tensor& col_seq_tensor = context->input(3);

    const int64_t out_batch = out_backprop.dim_size(0);
    const int64_t out_rows = out_backprop.dim_size(1);
    const int64_t out_cols = out_backprop.dim_size(2);
    const int64_t out_depth = out_backprop.dim_size(3);

    OP_REQUIRES(context, row_seq_tensor.NumElements() > out_rows,
                errors::InvalidArgument("Given out_backprop shape ",
                                        out_backprop.shape().DebugString(),
                                        ", row_seq_tensor must have at least ",
                                        out_rows + 1, " elements, but got ",
                                        row_seq_tensor.NumElements()));
    OP_REQUIRES(context, col_seq_tensor.NumElements() > out_cols,
                errors::InvalidArgument("Given out_backprop shape ",
                                        out_backprop.shape().DebugString(),
                                        ", col_seq_tensor must have at least ",
                                        out_cols + 1, " elements, but got ",
                                        col_seq_tensor.NumElements()));

    auto row_seq_tensor_flat = row_seq_tensor.flat<int64>();
    auto col_seq_tensor_flat = col_seq_tensor.flat<int64>();
    auto orig_input_tensor_shape_flat = orig_input_tensor_shape.flat<int64>();

    const int64_t in_batch = orig_input_tensor_shape_flat(0);
    const int64_t in_rows = orig_input_tensor_shape_flat(1);
    const int64_t in_cols = orig_input_tensor_shape_flat(2);
    const int64_t in_depth = orig_input_tensor_shape_flat(3);

    constexpr int tensor_in_and_out_dims = 4;
    // Transform orig_input_tensor_shape into TensorShape
    TensorShape in_shape;
    for (auto i = 0; i < tensor_in_and_out_dims; ++i) {
      in_shape.AddDim(orig_input_tensor_shape_flat(i));
    }

    // Create intermediate in_backprop.
    Tensor in_backprop_tensor_temp;
    OP_REQUIRES_OK(context, context->forward_input_or_allocate_temp(
                                {0}, DataTypeToEnum<double>::v(), in_shape,
                                &in_backprop_tensor_temp));
    in_backprop_tensor_temp.flat<double>().setZero();
    // Transform 4D tensor to 2D matrix.
    EigenDoubleMatrixMap in_backprop_tensor_temp_mat(
        in_backprop_tensor_temp.flat<double>().data(), in_depth,
        in_cols * in_rows * in_batch);
    ConstEigenMatrixMap out_backprop_mat(out_backprop.flat<T>().data(),
                                         out_depth,
                                         out_cols * out_rows * out_batch);
    // Loop through each element of out_backprop and evenly distribute the
    // element to the corresponding pooling cell.
    const int64_t in_max_row_index = in_rows - 1;
    const int64_t in_max_col_index = in_cols - 1;
    for (int64_t b = 0; b < out_batch; ++b) {
      for (int64_t r = 0; r < out_rows; ++r) {
        const int64_t in_row_start = row_seq_tensor_flat(r);
        int64_t in_row_end = overlapping_ ? row_seq_tensor_flat(r + 1)
                                          : row_seq_tensor_flat(r + 1) - 1;
        in_row_end = std::min(in_row_end, in_max_row_index);
        for (int64_t c = 0; c < out_cols; ++c) {
          const int64_t in_col_start = col_seq_tensor_flat(c);
          int64_t in_col_end = overlapping_ ? col_seq_tensor_flat(c + 1)
                                            : col_seq_tensor_flat(c + 1) - 1;
          in_col_end = std::min(in_col_end, in_max_col_index);

          const int64_t num_elements_in_pooling_cell =
              (in_row_end - in_row_start + 1) * (in_col_end - in_col_start + 1);
          const int64_t out_index = (b * out_rows + r) * out_cols + c;
          // Now we can evenly distribute out_backprop(b, h, w, *) to
          // in_backprop(b, hs:he, ws:we, *).
          for (int64_t in_r = in_row_start; in_r <= in_row_end; ++in_r) {
            for (int64_t in_c = in_col_start; in_c <= in_col_end; ++in_c) {
              const int64_t in_index = (b * in_rows + in_r) * in_cols + in_c;
              // Walk through each channel (depth).
              for (int64_t d = 0; d < out_depth; ++d) {
                const double out_backprop_element = static_cast<double>(
                    out_backprop_mat.coeffRef(d, out_index));
                double& in_backprop_ref =
                    in_backprop_tensor_temp_mat.coeffRef(d, in_index);
                in_backprop_ref +=
                    out_backprop_element / num_elements_in_pooling_cell;
              }
            }
          }
        }
      }
    }

    // Depending on the type, cast double to type T.
    Tensor* in_backprop_tensor = nullptr;
    OP_REQUIRES_OK(context, context->forward_input_or_allocate_output(
                                {0}, 0, in_shape, &in_backprop_tensor));
    auto in_backprop_tensor_flat = in_backprop_tensor->flat<T>();
    auto in_backprop_tensor_temp_flat = in_backprop_tensor_temp.flat<double>();
    for (int64_t i = 0; i < in_backprop_tensor_flat.size(); ++i) {
      in_backprop_tensor_flat(i) =
          static_cast<T>(in_backprop_tensor_temp_flat(i));
    }
  }

Validate dimensions of input tensor in `FractionalAvgPoolGrad`

PiperOrigin-RevId: 388286227
Change-Id: Ieb7566155e92acc8993a2212c76deacadc0edc8a


  void Compute(OpKernelContext* context) override {
    // Here's the basic idea:
    // Batch and depth dimension are independent from row and col dimension. And
    // because FractionalAvgPool currently only support pooling along row and
    // col, we can basically think of this 4D tensor backpropagation as
    // operation of a series of 2D planes.
    //
    // For each element of a 'slice' (2D plane) of output_backprop, we need to
    // figure out its contributors when doing FractionalAvgPool operation. This
    // can be done based on row_pooling_sequence, col_pooling_seq and
    // overlapping.
    // Once we figure out the original contributors, we just need to evenly
    // divide the value of this element among these contributors.
    //
    // Internally, we divide the out_backprop tensor and store it in a temporary
    // tensor of double type. And cast it to the corresponding type.
    typedef Eigen::Map<const Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>>
        ConstEigenMatrixMap;
    typedef Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>>
        EigenDoubleMatrixMap;

    // Grab the inputs.
    const Tensor& orig_input_tensor_shape = context->input(0);
    OP_REQUIRES(context,
                orig_input_tensor_shape.dims() == 1 &&
                    orig_input_tensor_shape.NumElements() == 4,
                errors::InvalidArgument("original input tensor shape must be"
                                        "1-dimensional and 4 elements"));
    const Tensor& out_backprop = context->input(1);
    const Tensor& row_seq_tensor = context->input(2);
    const Tensor& col_seq_tensor = context->input(3);

    const int64_t out_batch = out_backprop.dim_size(0);
    const int64_t out_rows = out_backprop.dim_size(1);
    const int64_t out_cols = out_backprop.dim_size(2);
    const int64_t out_depth = out_backprop.dim_size(3);

    OP_REQUIRES(context, row_seq_tensor.NumElements() > out_rows,
                errors::InvalidArgument("Given out_backprop shape ",
                                        out_backprop.shape().DebugString(),
                                        ", row_seq_tensor must have at least ",
                                        out_rows + 1, " elements, but got ",
                                        row_seq_tensor.NumElements()));
    OP_REQUIRES(context, col_seq_tensor.NumElements() > out_cols,
                errors::InvalidArgument("Given out_backprop shape ",
                                        out_backprop.shape().DebugString(),
                                        ", col_seq_tensor must have at least ",
                                        out_cols + 1, " elements, but got ",
                                        col_seq_tensor.NumElements()));

    auto row_seq_tensor_flat = row_seq_tensor.flat<int64>();
    auto col_seq_tensor_flat = col_seq_tensor.flat<int64>();
    auto orig_input_tensor_shape_flat = orig_input_tensor_shape.flat<int64>();

    const int64_t in_batch = orig_input_tensor_shape_flat(0);
    const int64_t in_rows = orig_input_tensor_shape_flat(1);
    const int64_t in_cols = orig_input_tensor_shape_flat(2);
    const int64_t in_depth = orig_input_tensor_shape_flat(3);
    OP_REQUIRES(
        context, in_batch != 0,
        errors::InvalidArgument("Batch dimension of input must not be 0"));
    OP_REQUIRES(
        context, in_rows != 0,
        errors::InvalidArgument("Rows dimension of input must not be 0"));
    OP_REQUIRES(
        context, in_cols != 0,
        errors::InvalidArgument("Columns dimension of input must not be 0"));
    OP_REQUIRES(
        context, in_depth != 0,
        errors::InvalidArgument("Depth dimension of input must not be 0"));

    constexpr int tensor_in_and_out_dims = 4;
    // Transform orig_input_tensor_shape into TensorShape
    TensorShape in_shape;
    for (auto i = 0; i < tensor_in_and_out_dims; ++i) {
      in_shape.AddDim(orig_input_tensor_shape_flat(i));
    }

    // Create intermediate in_backprop.
    Tensor in_backprop_tensor_temp;
    OP_REQUIRES_OK(context, context->forward_input_or_allocate_temp(
                                {0}, DataTypeToEnum<double>::v(), in_shape,
                                &in_backprop_tensor_temp));
    in_backprop_tensor_temp.flat<double>().setZero();
    // Transform 4D tensor to 2D matrix.
    EigenDoubleMatrixMap in_backprop_tensor_temp_mat(
        in_backprop_tensor_temp.flat<double>().data(), in_depth,
        in_cols * in_rows * in_batch);
    ConstEigenMatrixMap out_backprop_mat(out_backprop.flat<T>().data(),
                                         out_depth,
                                         out_cols * out_rows * out_batch);
    // Loop through each element of out_backprop and evenly distribute the
    // element to the corresponding pooling cell.
    const int64_t in_max_row_index = in_rows - 1;
    const int64_t in_max_col_index = in_cols - 1;
    for (int64_t b = 0; b < out_batch; ++b) {
      for (int64_t r = 0; r < out_rows; ++r) {
        const int64_t in_row_start = row_seq_tensor_flat(r);
        int64_t in_row_end = overlapping_ ? row_seq_tensor_flat(r + 1)
                                          : row_seq_tensor_flat(r + 1) - 1;
        in_row_end = std::min(in_row_end, in_max_row_index);
        for (int64_t c = 0; c < out_cols; ++c) {
          const int64_t in_col_start = col_seq_tensor_flat(c);
          int64_t in_col_end = overlapping_ ? col_seq_tensor_flat(c + 1)
                                            : col_seq_tensor_flat(c + 1) - 1;
          in_col_end = std::min(in_col_end, in_max_col_index);

          const int64_t num_elements_in_pooling_cell =
              (in_row_end - in_row_start + 1) * (in_col_end - in_col_start + 1);
          const int64_t out_index = (b * out_rows + r) * out_cols + c;
          // Now we can evenly distribute out_backprop(b, h, w, *) to
          // in_backprop(b, hs:he, ws:we, *).
          for (int64_t in_r = in_row_start; in_r <= in_row_end; ++in_r) {
            for (int64_t in_c = in_col_start; in_c <= in_col_end; ++in_c) {
              const int64_t in_index = (b * in_rows + in_r) * in_cols + in_c;
              // Walk through each channel (depth).
              for (int64_t d = 0; d < out_depth; ++d) {
                const double out_backprop_element = static_cast<double>(
                    out_backprop_mat.coeffRef(d, out_index));
                double& in_backprop_ref =
                    in_backprop_tensor_temp_mat.coeffRef(d, in_index);
                in_backprop_ref +=
                    out_backprop_element / num_elements_in_pooling_cell;
              }
            }
          }
        }
      }
    }

    // Depending on the type, cast double to type T.
    Tensor* in_backprop_tensor = nullptr;
    OP_REQUIRES_OK(context, context->forward_input_or_allocate_output(
                                {0}, 0, in_shape, &in_backprop_tensor));
    auto in_backprop_tensor_flat = in_backprop_tensor->flat<T>();
    auto in_backprop_tensor_temp_flat = in_backprop_tensor_temp.flat<double>();
    for (int64_t i = 0; i < in_backprop_tensor_flat.size(); ++i) {
      in_backprop_tensor_flat(i) =
          static_cast<T>(in_backprop_tensor_temp_flat(i));
    }
  }