static GF_Err txtin_setup_ttxt(GF_Filter *filter, GF_TXTIn *ctx)
{
	GF_Err e;
	u32 j, k, ID, OCR_ES_ID;
	u64 file_size;
	GF_XMLNode *root, *ext;
	GF_PropertyValue *dcd;

	ctx->parser = gf_xml_dom_new();
	e = gf_xml_dom_parse(ctx->parser, ctx->file_name, ttxt_dom_progress, ctx);
	if (e) {
		GF_LOG(GF_LOG_ERROR, GF_LOG_PARSER, ("[TXTIn] Error parsing TTXT file: Line %d - %s\n", gf_xml_dom_get_line(ctx->parser), gf_xml_dom_get_error(ctx->parser)));
		return e;
	}
	root = gf_xml_dom_get_root(ctx->parser);

	if (strcmp(root->name, "TextStream")) {
		GF_LOG(GF_LOG_ERROR, GF_LOG_PARSER, ("[TXTIn] Invalid Timed Text file - expecting \"TextStream\" got %s", root->name));
		return GF_NON_COMPLIANT_BITSTREAM;
	}
	file_size = ctx->end;
	ctx->end = 0;

	/*setup track in 3GP format directly (no ES desc)*/
	if (!ctx->timescale) ctx->timescale = 1000;
	OCR_ES_ID = ID = 0;

	if (!ctx->opid) ctx->opid = gf_filter_pid_new(filter);
	gf_filter_pid_set_property(ctx->opid, GF_PROP_PID_STREAM_TYPE, &PROP_UINT(GF_STREAM_TEXT) );
	gf_filter_pid_set_property(ctx->opid, GF_PROP_PID_CODECID, &PROP_UINT(GF_CODECID_TX3G) );
	gf_filter_pid_set_property(ctx->opid, GF_PROP_PID_TIMESCALE, &PROP_UINT(ctx->timescale) );
	gf_filter_pid_set_property(ctx->opid, GF_PROP_PID_DOWN_SIZE, &PROP_LONGUINT(file_size) );

	if (!ID) ID = 1;
	gf_filter_pid_set_property(ctx->opid, GF_PROP_PID_ID, &PROP_UINT(ID) );
	if (OCR_ES_ID) gf_filter_pid_set_property(ctx->opid, GF_PROP_PID_CLOCK_ID, &PROP_UINT(OCR_ES_ID) );

	ctx->nb_children = gf_list_count(root->content);

	ctx->cur_child_idx = 0;
	for (ctx->cur_child_idx=0; ctx->cur_child_idx < ctx->nb_children; ctx->cur_child_idx++) {
		GF_XMLNode *node = (GF_XMLNode*) gf_list_get(root->content, ctx->cur_child_idx);

		if (node->type) {
			continue;
		}

		if (!strcmp(node->name, "TextStreamHeader")) {
			GF_XMLNode *sdesc;
			s32 w, h, tx, ty, layer;
			u32 tref_id;
			GF_XMLAttribute *att;
			w = ctx->width;
			h = ctx->height;
			tx = ctx->txtx;
			ty = ctx->txty;
			layer = ctx->zorder;
			tref_id = 0;

			j=0;
			while ( (att=(GF_XMLAttribute *)gf_list_enum(node->attributes, &j))) {
				if (!strcmp(att->name, "width")) w = atoi(att->value);
				else if (!strcmp(att->name, "height")) h = atoi(att->value);
				else if (!strcmp(att->name, "layer")) layer = atoi(att->value);
				else if (!strcmp(att->name, "translation_x")) tx = atoi(att->value);
				else if (!strcmp(att->name, "translation_y")) ty = atoi(att->value);
				else if (!strcmp(att->name, "trefID")) tref_id = atoi(att->value);
			}

			if (tref_id) {
				gf_filter_pid_set_property_str(ctx->opid, "tref:chap", &PROP_UINT(tref_id) );
			}

			if (w) gf_filter_pid_set_property(ctx->opid, GF_PROP_PID_WIDTH, &PROP_UINT(w) );
			if (h) gf_filter_pid_set_property(ctx->opid, GF_PROP_PID_HEIGHT, &PROP_UINT(h) );
			if (tx) gf_filter_pid_set_property(ctx->opid, GF_PROP_PID_TRANS_X, &PROP_UINT(tx) );
			if (ty) gf_filter_pid_set_property(ctx->opid, GF_PROP_PID_TRANS_X, &PROP_UINT(ty) );
			if (layer) gf_filter_pid_set_property(ctx->opid, GF_PROP_PID_ZORDER, &PROP_SINT(ctx->zorder) );
			if (ctx->lang) gf_filter_pid_set_property(ctx->opid, GF_PROP_PID_LANGUAGE, &PROP_STRING( ctx->lang) );

			j=0;
			while ( (sdesc=(GF_XMLNode*)gf_list_enum(node->content, &j))) {
				if (sdesc->type) continue;

				if (!strcmp(sdesc->name, "TextSampleDescription")) {
					GF_TextSampleDescriptor td;
					memset(&td, 0, sizeof(GF_TextSampleDescriptor));
					td.tag = GF_ODF_TEXT_CFG_TAG;
					td.vert_justif = (s8) -1;
					td.default_style.fontID = 1;
					td.default_style.font_size = ctx->fontsize;

					k=0;
					while ( (att=(GF_XMLAttribute *)gf_list_enum(sdesc->attributes, &k))) {
						if (!strcmp(att->name, "horizontalJustification")) {
							if (!stricmp(att->value, "center")) td.horiz_justif = 1;
							else if (!stricmp(att->value, "right")) td.horiz_justif = (s8) -1;
							else if (!stricmp(att->value, "left")) td.horiz_justif = 0;
						}
						else if (!strcmp(att->name, "verticalJustification")) {
							if (!stricmp(att->value, "center")) td.vert_justif = 1;
							else if (!stricmp(att->value, "bottom")) td.vert_justif = (s8) -1;
							else if (!stricmp(att->value, "top")) td.vert_justif = 0;
						}
						else if (!strcmp(att->name, "backColor")) td.back_color = ttxt_get_color(att->value);
						else if (!strcmp(att->name, "verticalText") && !stricmp(att->value, "yes") ) td.displayFlags |= GF_TXT_VERTICAL;
						else if (!strcmp(att->name, "fillTextRegion") && !stricmp(att->value, "yes") ) td.displayFlags |= GF_TXT_FILL_REGION;
						else if (!strcmp(att->name, "continuousKaraoke") && !stricmp(att->value, "yes") ) td.displayFlags |= GF_TXT_KARAOKE;
						else if (!strcmp(att->name, "scroll")) {
							if (!stricmp(att->value, "inout")) td.displayFlags |= GF_TXT_SCROLL_IN | GF_TXT_SCROLL_OUT;
							else if (!stricmp(att->value, "in")) td.displayFlags |= GF_TXT_SCROLL_IN;
							else if (!stricmp(att->value, "out")) td.displayFlags |= GF_TXT_SCROLL_OUT;
						}
						else if (!strcmp(att->name, "scrollMode")) {
							u32 scroll_mode = GF_TXT_SCROLL_CREDITS;
							if (!stricmp(att->value, "Credits")) scroll_mode = GF_TXT_SCROLL_CREDITS;
							else if (!stricmp(att->value, "Marquee")) scroll_mode = GF_TXT_SCROLL_MARQUEE;
							else if (!stricmp(att->value, "Right")) scroll_mode = GF_TXT_SCROLL_RIGHT;
							else if (!stricmp(att->value, "Down")) scroll_mode = GF_TXT_SCROLL_DOWN;
							td.displayFlags |= ((scroll_mode<<7) & GF_TXT_SCROLL_DIRECTION);
						}
					}

					k=0;
					while ( (ext=(GF_XMLNode*)gf_list_enum(sdesc->content, &k))) {
						if (ext->type) continue;
						if (!strcmp(ext->name, "TextBox")) ttxt_parse_text_box(ext, &td.default_pos);
						else if (!strcmp(ext->name, "Style")) ttxt_parse_text_style(ctx, ext, &td.default_style);
						else if (!strcmp(ext->name, "FontTable")) {
							GF_XMLNode *ftable;
							u32 z=0;
							while ( (ftable=(GF_XMLNode*)gf_list_enum(ext->content, &z))) {
								u32 m;
								if (ftable->type || strcmp(ftable->name, "FontTableEntry")) continue;
								td.font_count += 1;
								td.fonts = (GF_FontRecord*)gf_realloc(td.fonts, sizeof(GF_FontRecord)*td.font_count);
								m=0;
								while ( (att=(GF_XMLAttribute *)gf_list_enum(ftable->attributes, &m))) {
									if (!stricmp(att->name, "fontID")) td.fonts[td.font_count-1].fontID = atoi(att->value);
									else if (!stricmp(att->name, "fontName")) td.fonts[td.font_count-1].fontName = gf_strdup(att->value);
								}
							}
						}
					}
					if (ctx->nodefbox) {
						td.default_pos.top = td.default_pos.left = td.default_pos.right = td.default_pos.bottom = 0;
					} else {
						if ((td.default_pos.bottom==td.default_pos.top) || (td.default_pos.right==td.default_pos.left)) {
							td.default_pos.top = td.default_pos.left = 0;
							td.default_pos.right = w;
							td.default_pos.bottom = h;
						}
					}
					if (!td.fonts) {
						td.font_count = 1;
						td.fonts = (GF_FontRecord*)gf_malloc(sizeof(GF_FontRecord));
						td.fonts[0].fontID = 1;
						td.fonts[0].fontName = gf_strdup("Serif");
					}
					GF_SAFEALLOC(dcd, GF_PropertyValue);
					if (dcd) {
						dcd->type = GF_PROP_DATA;

						gf_odf_tx3g_write(&td, &dcd->value.data.ptr, &dcd->value.data.size);
						if (!ctx->text_descs) ctx->text_descs = gf_list_new();
						gf_list_add(ctx->text_descs, dcd);
					}

					for (k=0; k<td.font_count; k++) gf_free(td.fonts[k].fontName);
					gf_free(td.fonts);
				}
			}
		}
		else {
			break;
		}
	}

	if (!ctx->text_descs) {
		GF_LOG(GF_LOG_ERROR, GF_LOG_PARSER, ("[TXTIn] Invalid Timed Text file - text stream header not found or empty\n"));
		return GF_NON_COMPLIANT_BITSTREAM;
	}
	dcd = gf_list_get(ctx->text_descs, 0);
	gf_filter_pid_set_property(ctx->opid, GF_PROP_PID_DECODER_CONFIG, dcd);
	ctx->last_desc_idx = 1;

	ctx->first_samp = GF_TRUE;
	ctx->last_sample_empty = GF_FALSE;
	ctx->last_sample_duration = 0;

	txtin_probe_duration(ctx);

	return GF_OK;
}

fixed #1897


static GF_Err swf_svg_add_iso_header(void *user, const u8 *data, u32 length, Bool isHeader)
{
	GF_TXTIn *ctx = (GF_TXTIn *)user;

	if (isHeader) {
		if (!ctx->hdr_parsed) {
			gf_filter_pid_set_property(ctx->opid, GF_PROP_PID_DECODER_CONFIG, &PROP_DATA((char *)data, (u32) ( strlen(data)+1 ) )  );
			ctx->hdr_parsed = GF_TRUE;
		}
	} else if (!ctx->seek_state) {
		GF_FilterPacket *pck;
		u8 *pck_data;
		pck = gf_filter_pck_new_alloc(ctx->opid, length, &pck_data);
		if (pck) {
			memcpy(pck_data, data, length);
			gf_filter_pck_set_framing(pck, GF_FALSE, GF_TRUE);

			gf_filter_pck_send(pck);
		}
	}
	return GF_OK;
}

fixed #1897


static GF_Err gf_text_process_ttml(GF_Filter *filter, GF_TXTIn *ctx)
{
	GF_Err e;
	GF_XMLNode *root;
	u32 i, nb_res_interval=0, k, nb_div_nodes;
	char *samp_text=NULL;
	GF_List *emb_resources = NULL;
	TTMLInterval *interval;
	Bool sample_empty = GF_TRUE;

	if (!ctx->is_setup) return gf_text_ttml_setup(filter, ctx);
	if (ctx->non_compliant_ttml || !ctx->opid) return GF_NOT_SUPPORTED;
	if (!ctx->playstate) return GF_OK;
	else if (ctx->playstate==2) return GF_EOS;

	if (ctx->seek_state==1) {
		ctx->seek_state = 2;
		ctx->current_tt_interval = 0;
	}

	interval = gf_list_get(ctx->intervals, ctx->current_tt_interval);
	if (!interval) {
		GF_LOG(GF_LOG_DEBUG, GF_LOG_PARSER, ("[TTML EBU-TTD] last_sample_duration="LLU", last_sample_end="LLU"\n", ctx->last_sample_duration, ctx->end));

		gf_filter_pid_set_info_str( ctx->opid, "ttxt:last_dur", &PROP_UINT((u32) ctx->last_sample_duration) );
		gf_filter_pid_set_eos(ctx->opid);
		return GF_EOS;
	}
	ctx->current_tt_interval++;

	emb_resources = interval->resources ? interval->resources : ctx->ttml_resources;
	nb_res_interval = gf_list_count(emb_resources);

	root = gf_xml_dom_get_root(ctx->parser);

	nb_div_nodes = gf_list_count(ctx->div_nodes_list);
	for (k=0; k<nb_div_nodes; k++) {
		Bool has_content = GF_FALSE;
		GF_XMLNode *div_node = gf_list_get(ctx->div_nodes_list, k);
		u32 nb_children = gf_list_count(div_node->content);

		GF_XMLNode *copy_div_node = gf_list_get(ctx->body_node->content, k);

		for (i=0; i < nb_children; i++) {
			GF_XMLNode *p_node;
			GF_XMLAttribute *p_att;
			u32 p_idx = 0;
			s64 ts_begin = -1, ts_end = -1;
			Bool in_range;
			GF_XMLNode *div_child = (GF_XMLNode*)gf_list_get(div_node->content, i);
			if (div_child->type) {
				continue;
			}
			e = gf_xml_get_element_check_namespace(div_child, "p", root->ns);
			if (e == GF_BAD_PARAM) {
				GF_LOG(GF_LOG_WARNING, GF_LOG_PARSER, ("[TTML EBU-TTD] ignored \"%s\" node, check your namespaces\n", div_child->name));
				continue;
			}

			//sample is either in the <p> ...
			p_idx = 0;
			while ( (p_att = (GF_XMLAttribute*)gf_list_enum(div_child->attributes, &p_idx))) {
				if (!strcmp(p_att->name, "begin")) {
					if (ts_begin != -1) {
						GF_LOG(GF_LOG_ERROR, GF_LOG_PARSER, ("[TTML EBU-TTD] duplicated \"begin\" attribute. Abort.\n"));
						e = GF_NON_COMPLIANT_BITSTREAM;
						goto exit;
					}
					ts_begin = ttml_get_timestamp(ctx, p_att->value);
				} else if (!strcmp(p_att->name, "end")) {
					if (ts_end != -1) {
						GF_LOG(GF_LOG_ERROR, GF_LOG_PARSER, ("[TTML EBU-TTD] duplicated \"end\" attribute. Abort.\n"));
						e = GF_NON_COMPLIANT_BITSTREAM;
						goto exit;
					}
					ts_end = ttml_get_timestamp(ctx, p_att->value);
				}
			}

			in_range = ttml_check_range(interval, ts_begin, ts_end);
			if (in_range) {
				GF_XMLNode *prev_child = i ? (GF_XMLNode*) gf_list_get(div_node->content, i-1) : NULL;
				if (prev_child && prev_child->type) {
					gf_xml_dom_append_child(copy_div_node, prev_child);
				}
				e = gf_xml_dom_append_child(copy_div_node, div_child);
				assert(e == GF_OK);
				has_content = GF_TRUE;
			}

			//or under a <span>
			p_idx = 0;
			while ( (p_node = (GF_XMLNode*)gf_list_enum(div_child->content, &p_idx))) {
				u32 span_idx = 0;
				GF_XMLAttribute *span_att;
				e = gf_xml_get_element_check_namespace(p_node, "span", root->ns);
				if (e == GF_BAD_PARAM) {
					GF_LOG(GF_LOG_WARNING, GF_LOG_PARSER, ("[TTML EBU-TTD] ignored \"%s\" node, check your namespaces\n", p_node->name));
				}
				else if (e)
					continue;

				ts_begin = ts_end = -1;
				while ( (span_att = (GF_XMLAttribute*)gf_list_enum(p_node->attributes, &span_idx))) {
					if (!strcmp(span_att->name, "begin")) {
						if (ts_begin != -1) {
							GF_LOG(GF_LOG_ERROR, GF_LOG_PARSER, ("[TTML EBU-TTD] duplicated \"begin\" attribute under <span>. Abort.\n"));
							e = GF_NON_COMPLIANT_BITSTREAM;
							goto exit;
						}
						ts_begin = ttml_get_timestamp(ctx, span_att->value);
					} else if (!strcmp(span_att->name, "end")) {
						if (ts_end != -1) {
							GF_LOG(GF_LOG_ERROR, GF_LOG_PARSER, ("[TTML EBU-TTD] duplicated \"end\" attribute under <span>. Abort.\n"));
							e = GF_NON_COMPLIANT_BITSTREAM;
							goto exit;
						}
						ts_end = ttml_get_timestamp(ctx, span_att->value);
					}
				}

				/*append the entire <p> and break (we cannot split the text content)*/
				in_range = ttml_check_range(interval, ts_begin, ts_end);
				if (in_range) {
					GF_XMLNode *prev_child = i ? (GF_XMLNode*) gf_list_get(div_node->content, i-1) : NULL;
					if (prev_child && prev_child->type) {
						gf_xml_dom_append_child(copy_div_node, prev_child);
					}
					e = gf_xml_dom_append_child(copy_div_node, div_child);
					assert(e == GF_OK);
					has_content = GF_TRUE;
					break;
				}
			}
		}
		if (has_content) {
			GF_XMLNode *last_child = (GF_XMLNode*) gf_list_last(div_node->content);
			if (last_child && last_child->type) {
				gf_xml_dom_append_child(copy_div_node, last_child);
			}
			sample_empty = GF_FALSE;
		}
	}

	//empty doc
	if (!ctx->body_node)
		sample_empty = GF_FALSE;

	if (! sample_empty) {
		samp_text = gf_xml_dom_serialize_root((GF_XMLNode*)ctx->root_working_copy, GF_FALSE, GF_FALSE);

		for (k=0; k<nb_div_nodes; k++) {
			GF_XMLNode *copy_div_node = gf_list_get(ctx->body_node->content, k);
			if (!copy_div_node->type)
				gf_list_reset(copy_div_node->content);
		}
	}

	if (samp_text) {
		GF_FilterPacket *pck;
		u8 *pck_data;
		Bool skip_pck = GF_FALSE;
		u32 txt_len;
		u32 res_len = 0;
		char *txt_str;

		if (interval->begin < (s64) ctx->end) {
			GF_LOG(GF_LOG_ERROR, GF_LOG_PARSER, ("[TTML EBU-TTD] Error computing overlapped intervals! \"begin\" is "LLD" , last \"end\" was "LLD". Abort.\n", interval->begin, ctx->end));
			e = GF_NOT_SUPPORTED;
			goto exit;
		}

		txt_str = ttxt_parse_string(samp_text, GF_TRUE);
		if (!txt_str) txt_str = "";
		txt_len = (u32) strlen(txt_str);

		for (i=0; i<nb_res_interval; i++) {
			TTMLRes *res = gf_list_get(emb_resources, i);
			res_len += res->size;
		}

		if (ctx->first_samp) {
			interval->begin = 0; /*in MP4 we must start at T=0*/
			ctx->first_samp = GF_FALSE;
		}

		ctx->last_sample_duration = interval->end - interval->begin;

		ctx->end = interval->end;
		GF_LOG(GF_LOG_DEBUG, GF_LOG_PARSER, ("ts_begin="LLD", ts_end="LLD", last_sample_duration="LLU" (real duration: "LLU"), last_sample_end="LLU"\n", interval->begin, interval->end, interval->end - ctx->end, ctx->last_sample_duration, ctx->end));

		if (ctx->seek_state==2) {
			Double end = (Double) interval->end;
			end /= ctx->timescale;
			if (end<ctx->start_range) skip_pck = GF_TRUE;
			else ctx->seek_state = 0;
		}

		if (!skip_pck) {
			pck = gf_filter_pck_new_alloc(ctx->opid, txt_len+res_len, &pck_data);
			if (!pck) {
				gf_free(samp_text);
				return GF_OUT_OF_MEM;
			}
			memcpy(pck_data, txt_str, txt_len);
			gf_filter_pck_set_sap(pck, GF_FILTER_SAP_1);

			if (ctx->ttml_dur>0) {
				gf_filter_pck_set_cts(pck, 0);
				gf_filter_pck_set_duration(pck, (u32) ctx->ttml_dur);
				ctx->last_sample_duration = (u64) ctx->ttml_dur * 1000 / ctx->timescale;
			} else {
				gf_filter_pck_set_cts(pck, (ctx->timescale * interval->begin)/1000);
				if (interval->end >= interval->begin) {
					gf_filter_pck_set_duration(pck, (u32) ((ctx->timescale * (interval->end - interval->begin) )/1000) );
				}
			}

			if (res_len) {
				GF_BitStream *subs = gf_bs_new(NULL, 0, GF_BITSTREAM_WRITE);
				u8 *subs_data;
				u32 subs_size;
				//subs 0
				gf_bs_write_u32(subs, 0);
				gf_bs_write_u32(subs, txt_len);
				gf_bs_write_u32(subs, 0);
				gf_bs_write_u8(subs, 0);
				gf_bs_write_u8(subs, 0);

				pck_data += txt_len;
				for (i=0; i<nb_res_interval; i++) {
					TTMLRes *res = gf_list_get(emb_resources, i);
					memcpy(pck_data, res->data, res->size);
					pck_data += res->size;

					//subs >0
					gf_bs_write_u32(subs, 0);
					gf_bs_write_u32(subs, res->size);
					gf_bs_write_u32(subs, 0);
					gf_bs_write_u8(subs, 0);
					gf_bs_write_u8(subs, 0);
				}
				gf_bs_get_content(subs, &subs_data, &subs_size);
				gf_bs_del(subs);
				gf_filter_pck_set_property(pck, GF_PROP_PCK_SUBS, &PROP_DATA_NO_COPY(subs_data, subs_size) );
			}
			gf_filter_pck_send(pck);
		}

		gf_free(samp_text);
		samp_text = NULL;
	} else {
		GF_LOG(GF_LOG_WARNING, GF_LOG_PARSER, ("[TTML EBU-TTD] empty sample (begin="LLD", end="LLD"). Skip.\n", interval->begin, interval->end));
	}

	return GF_OK;

exit:
	if (!ctx->non_compliant_ttml) {
		ctx->non_compliant_ttml = GF_TRUE;
		gf_filter_pid_set_discard(ctx->ipid, GF_TRUE);
	}
	return e;
}

fixed #1897


static void gf_webvtt_flush_sample(void *user, GF_WebVTTSample *samp)
{
	u64 start, end;
	GF_TXTIn *ctx = (GF_TXTIn *)user;
	GF_ISOSample *s;

	start = gf_webvtt_sample_get_start(samp);
	end = gf_webvtt_sample_get_end(samp);

	if (ctx->seek_state==2) {
		Double tsend = (Double) end;
		tsend /= 1000;
		if (tsend<ctx->start_range) return;
		ctx->seek_state = 0;
	}

	s = gf_isom_webvtt_to_sample(samp);
	if (s) {
		GF_FilterPacket *pck;
		u8 *pck_data;

		pck = gf_filter_pck_new_alloc(ctx->opid, s->dataLength, &pck_data);
		if (pck) {
			memcpy(pck_data, s->data, s->dataLength);
			gf_filter_pck_set_cts(pck, (u64) (ctx->timescale * start / 1000) );
			gf_filter_pck_set_sap(pck, GF_FILTER_SAP_1);

			if (end && (end>=start) ) {
				gf_filter_pck_set_duration(pck, (u32) (ctx->timescale * (end-start) / 1000) );
			}
			gf_filter_pck_send(pck);
		}

		gf_isom_sample_del(&s);
	}
	gf_webvtt_sample_del(samp);

	gf_filter_pid_set_info(ctx->opid, GF_PROP_PID_DOWN_BYTES, &PROP_LONGUINT( gf_ftell(ctx->src )) );

	if (gf_filter_pid_would_block(ctx->opid))
		gf_webvtt_parser_suspend(ctx->vttparser);

}

fixed #1897


static GF_Err ttml_setup_intervals(GF_TXTIn *ctx)
{
	u32 k, i, nb_divs;
	s64 ttml_zero_ms = 0;
	GF_Err e;
	GF_XMLNode *root;

	if (!ctx->intervals)
		ctx->intervals = gf_list_new();
	else
		ttml_reset_intervals(ctx);

	ctx->has_images = GF_FALSE;
	root = ctx->root_working_copy;
	for (k=0; k<gf_list_count(root->content); k++) {
		GF_XMLNode *head = (GF_XMLNode*)gf_list_get(root->content, k);
		if (head->type) continue;
		if (strcmp(head->name, "head")) continue;
		ttml_push_resources(ctx, NULL, head, NULL);
		break;
	}

	root = gf_xml_dom_get_root(ctx->parser);
	if (ctx->ttml_zero) {
		if (ctx->ttml_zero[0]=='T')
			ttml_zero_ms = ttml_get_timestamp(ctx, (char *) ctx->ttml_zero+1);
		else
			ttml_zero_ms = ttml_get_timestamp(ctx, (char *) ctx->ttml_zero);
	}

	nb_divs = gf_list_count(ctx->div_nodes_list);
	for (i=0; i<nb_divs; i++) {
		u32 nb_children;
		GF_XMLNode *div_node = gf_list_get(ctx->div_nodes_list, i);
		nb_children = gf_list_count(div_node->content);

		for (k=0; k<nb_children; k++) {
			TTMLInterval *ival=NULL;
			u32 p_idx;
			Bool drop = GF_FALSE;
			GF_XMLAttribute *p_att;
			GF_XMLNode *p_node;
			s64 begin=-1, end=-1;
			GF_XMLNode *adiv_child = (GF_XMLNode*)gf_list_get(div_node->content, k);
			if (adiv_child->type) continue;
			e = gf_xml_get_element_check_namespace(adiv_child, "p", root->ns);
			if (e) continue;

			p_idx = 0;
			while ( (p_att = (GF_XMLAttribute*)gf_list_enum(adiv_child->attributes, &p_idx))) {
				if (!strcmp(p_att->name, "begin")) {
					e = ttml_rewrite_timestamp(ctx, ttml_zero_ms, p_att, &begin, &drop);
					if (e) return e;
				}
				if (!strcmp(p_att->name, "end")) {
					e = ttml_rewrite_timestamp(ctx, ttml_zero_ms, p_att, &end, &drop);
					if (e) return e;
				}
			}
			if (drop) {
				gf_xml_dom_node_del(adiv_child);
				gf_list_rem(div_node->content, k);
				k--;
				nb_children--;
				continue;
			}

			e = ttml_push_interval(ctx, begin, end, &ival);
			if (e) return e;

			e = ttml_push_resources(ctx, ival, adiv_child, NULL);
			if (e) return e;

			p_idx = 0;
			while ( (p_node = (GF_XMLNode*)gf_list_enum(adiv_child->content, &p_idx))) {
				s64 s_begin=-1, s_end=-1;
				e = gf_xml_get_element_check_namespace(p_node, "span", root->ns);
				if (e) continue;

				u32 span_idx = 0;
				GF_XMLAttribute *span_att;
				while ( (span_att = (GF_XMLAttribute*)gf_list_enum(p_node->attributes, &span_idx))) {
					if (!strcmp(span_att->name, "begin")) {
						e = ttml_rewrite_timestamp(ctx, ttml_zero_ms, span_att, &s_begin, &drop);
						if (e) return e;
					} else if (!strcmp(span_att->name, "end")) {
						e = ttml_rewrite_timestamp(ctx, ttml_zero_ms, span_att, &s_end, &drop);
						if (e) return e;
					}
				}
				e = ttml_push_interval(ctx, s_begin, s_end, &ival);
				if (e) return e;

				e = ttml_push_resources(ctx, ival, p_node, NULL);
				if (e) return e;
			}
		}
	}

	//empty doc
	if (!gf_list_count(ctx->intervals)) {
		TTMLInterval *interval;
		GF_SAFEALLOC(interval, TTMLInterval);
		interval->begin = interval->end = 0;
		gf_list_add(ctx->intervals, interval);
	}

#ifndef GPAC_DISABLE_LOG
	if (gf_log_tool_level_on(GF_LOG_PARSER, GF_LOG_DEBUG)) {
		for (k=0; k<gf_list_count(ctx->intervals); k++) {
			TTMLInterval *ival = gf_list_get(ctx->intervals, k);
			GF_LOG(GF_LOG_DEBUG, GF_LOG_PARSER, ("[TTML EBU-TTD] Interval %d: "LLU"-"LLU"\n", k+1, ival->begin, ival->end));
		}
	}
#endif
	return GF_OK;
}

fixed #1897


static GF_Err ttml_push_resources(GF_TXTIn *ctx, TTMLInterval *interval, GF_XMLNode *node, GF_XMLNode *parent_source_node)
{
	u32 i;
	char szURN[1024];
	u8 *f_data;
	u32 f_size;
	u32 idx;
	GF_Err e;
	GF_XMLAttribute *att, *data_type = NULL;
	GF_XMLNode *child;
	Bool is_source = GF_FALSE;
	Bool is_data = GF_FALSE;
	Bool check_src = GF_FALSE;

	if (!ctx->ttml_embed)
		return GF_OK;

	if (!strcmp(node->name, "source")) {
		is_source = GF_TRUE;
		check_src = GF_TRUE;
	}
	else if (!strcmp(node->name, "data")) {
		is_data = parent_source_node ? GF_TRUE : GF_FALSE;
		check_src = GF_TRUE;
	}
	//we don't embed chunks
	else if (!strcmp(node->name, "chunk")) {
		return GF_OK;
	}
	else if (!strcmp(node->name, "audio") || !strcmp(node->name, "font") || !strcmp(node->name, "image")) {
		check_src = GF_TRUE;
	}

	if (check_src) {
		i = 0;
		while ( (att = (GF_XMLAttribute*)gf_list_enum(node->attributes, &i))) {
			char *url;
			if (!att->value) continue;
			if (is_data && !strcmp(att->name, "type")) {
				data_type = att;
				continue;
			}
			if (strcmp(att->name, "src")) continue;
			if (att->value[0]=='#') continue;

			if (!strncmp(att->value, "file://", 7)) {}
			else if (strstr(att->value, "://"))
				continue;

			url = gf_url_concatenate(ctx->file_name, att->value);
			//embed image
			e = gf_file_load_data(url, &f_data, &f_size);
			gf_free(url);
			if (e) return e;

			e = ttml_push_res(ctx, interval, f_data, f_size);
			if (e) return e;

			idx = gf_list_count(interval ? interval->resources : ctx->ttml_resources);
			gf_free(att->value);
			sprintf(szURN, "urn:mpeg:14496-30:%d", idx);
			att->value = gf_strdup(szURN);
			if (!att->value) return GF_OUT_OF_MEM;

			ctx->has_images = GF_TRUE;
		}
	}

	i = 0;
	while ( (child = (GF_XMLNode*) gf_list_enum(node->content, &i))) {
		if (child->type) {
			if (!is_data) continue;
			u8 *data = child->name;
			u32 ilen = (u32) strlen(data);
			f_size = 3*ilen/4;
			f_data = gf_malloc(sizeof(u8) * f_size);

			f_size = gf_base64_decode(data, ilen, f_data, f_size);

			e = ttml_push_res(ctx, interval, f_data, f_size);
			if (e) return e;

			idx = gf_list_count(interval ? interval->resources : ctx->ttml_resources);
			sprintf(szURN, "urn:mpeg:14496-30:%d", idx);

			GF_SAFEALLOC(att, GF_XMLAttribute)
			if (att) {
				att->name = gf_strdup("src");
				att->value = gf_strdup(szURN);
				gf_list_add(parent_source_node->attributes, att);
			}
			if (!att || !att->value || !att->name) return GF_OUT_OF_MEM;
			if (data_type) {
				gf_list_del_item(node->attributes, data_type);
				gf_list_add(parent_source_node->attributes, data_type);
			}
			gf_xml_dom_node_reset(parent_source_node, GF_FALSE, GF_TRUE);
			ctx->has_images = GF_TRUE;
			return GF_OK;
		}

		e = ttml_push_resources(ctx, interval, child, is_source ? node : NULL);
		if (e) return e;
	}
	return GF_OK;
}

fixed #1897


static GF_Err txtin_texml_setup(GF_Filter *filter, GF_TXTIn *ctx)
{
	GF_Err e;
	u32 ID, OCR_ES_ID, i;
	u64 file_size;
	GF_XMLAttribute *att;
	GF_XMLNode *root;

	ctx->parser = gf_xml_dom_new();
	e = gf_xml_dom_parse(ctx->parser, ctx->file_name, ttxt_dom_progress, ctx);
	if (e) {
		GF_LOG(GF_LOG_ERROR, GF_LOG_PARSER, ("[TXTIn] Error parsing TeXML file: Line %d - %s", gf_xml_dom_get_line(ctx->parser), gf_xml_dom_get_error(ctx->parser) ));
		gf_xml_dom_del(ctx->parser);
		ctx->parser = NULL;
		return e;
	}

	root = gf_xml_dom_get_root(ctx->parser);

	if (strcmp(root->name, "text3GTrack")) {
		GF_LOG(GF_LOG_ERROR, GF_LOG_PARSER, ("[TXTIn] Invalid QT TeXML file - expecting root \"text3GTrack\" got \"%s\"", root->name));
		return GF_NON_COMPLIANT_BITSTREAM;
	}
	file_size = ctx->end;
	ctx->txml_timescale = 600;

	i=0;
	while ( (att=(GF_XMLAttribute *)gf_list_enum(root->attributes, &i))) {
		if (!strcmp(att->name, "trackWidth")) ctx->width = atoi(att->value);
		else if (!strcmp(att->name, "trackHeight")) ctx->height = atoi(att->value);
		else if (!strcmp(att->name, "layer")) ctx->zorder = atoi(att->value);
		else if (!strcmp(att->name, "timeScale")) ctx->txml_timescale = atoi(att->value);
		else if (!strcmp(att->name, "transform")) {
			Float fx, fy;
			sscanf(att->value, "translate(%f,%f)", &fx, &fy);
			ctx->txtx = (u32) fx;
			ctx->txty = (u32) fy;
		}
	}

	/*setup track in 3GP format directly (no ES desc)*/
	OCR_ES_ID = ID = 0;
	if (!ctx->timescale) ctx->timescale = 1000;

	if (!ctx->opid) ctx->opid = gf_filter_pid_new(filter);
	gf_filter_pid_set_property(ctx->opid, GF_PROP_PID_STREAM_TYPE, &PROP_UINT(GF_STREAM_TEXT) );
	gf_filter_pid_set_property(ctx->opid, GF_PROP_PID_CODECID, &PROP_UINT(GF_ISOM_SUBTYPE_TX3G) );
	gf_filter_pid_set_property(ctx->opid, GF_PROP_PID_TIMESCALE, &PROP_UINT(ctx->timescale) );
	gf_filter_pid_set_property(ctx->opid, GF_PROP_PID_DOWN_SIZE, &PROP_LONGUINT(file_size) );

	if (!ID) ID = 1;
	gf_filter_pid_set_property(ctx->opid, GF_PROP_PID_ID, &PROP_UINT(ID) );
	if (OCR_ES_ID) gf_filter_pid_set_property(ctx->opid, GF_PROP_PID_CLOCK_ID, &PROP_UINT(OCR_ES_ID) );
	if (ctx->width) gf_filter_pid_set_property(ctx->opid, GF_PROP_PID_WIDTH, &PROP_UINT(ctx->width) );
	if (ctx->height) gf_filter_pid_set_property(ctx->opid, GF_PROP_PID_HEIGHT, &PROP_UINT(ctx->height) );
	if (ctx->zorder) gf_filter_pid_set_property(ctx->opid, GF_PROP_PID_ZORDER, &PROP_SINT(ctx->zorder) );
	if (ctx->lang) gf_filter_pid_set_property(ctx->opid, GF_PROP_PID_LANGUAGE, &PROP_STRING( ctx->lang) );

	ctx->nb_children = gf_list_count(root->content);
	ctx->cur_child_idx = 0;
	txtin_probe_duration(ctx);

	return GF_OK;
}

fixed #1897


static void ttxt_parse_text_style(GF_TXTIn *ctx, GF_XMLNode *n, GF_StyleRecord *style)
{
	u32 i=0;
	GF_XMLAttribute *att;
	memset(style, 0, sizeof(GF_StyleRecord));
	style->fontID = 1;
	style->font_size = ctx->fontsize ;
	style->text_color = 0xFFFFFFFF;

	while ( (att=(GF_XMLAttribute *)gf_list_enum(n->attributes, &i))) {
		if (!stricmp(att->name, "fromChar")) style->startCharOffset = atoi(att->value);
		else if (!stricmp(att->name, "toChar")) style->endCharOffset = atoi(att->value);
		else if (!stricmp(att->name, "fontID")) style->fontID = atoi(att->value);
		else if (!stricmp(att->name, "fontSize")) style->font_size = atoi(att->value);
		else if (!stricmp(att->name, "color")) style->text_color = ttxt_get_color(att->value);
		else if (!stricmp(att->name, "styles")) {
			if (strstr(att->value, "Bold")) style->style_flags |= GF_TXT_STYLE_BOLD;
			if (strstr(att->value, "Italic")) style->style_flags |= GF_TXT_STYLE_ITALIC;
			if (strstr(att->value, "Underlined")) style->style_flags |= GF_TXT_STYLE_UNDERLINED;
			if (strstr(att->value, "Strikethrough")) style->style_flags |= GF_TXT_STYLE_STRIKETHROUGH;
		}
	}
}

fixed #1897


s32 gf_text_get_utf_type(FILE *in_src)
{
	u32 read;
	unsigned char BOM[5];
	read = (u32) gf_fread(BOM, 5, in_src);
	if ((s32) read < 1)
		return -1;

	if ((BOM[0]==0xFF) && (BOM[1]==0xFE)) {
		/*UTF32 not supported*/
		if (!BOM[2] && !BOM[3]) return -1;
		gf_fseek(in_src, 2, SEEK_SET);
		return 3;
	}
	if ((BOM[0]==0xFE) && (BOM[1]==0xFF)) {
		/*UTF32 not supported*/
		if (!BOM[2] && !BOM[3]) return -1;
		gf_fseek(in_src, 2, SEEK_SET);
		return 2;
	} else if ((BOM[0]==0xEF) && (BOM[1]==0xBB) && (BOM[2]==0xBF)) {
		gf_fseek(in_src, 3, SEEK_SET);
		return 1;
	}
	if (BOM[0]<0x80) {
		gf_fseek(in_src, 0, SEEK_SET);
		return 0;
	}
	return -1;
}

fixed #1897


static GF_Err gf_webvtt_import_report(void *user, GF_Err e, char *message, const char *line)
{
	GF_LOG(e ? GF_LOG_WARNING : GF_LOG_INFO, GF_LOG_AUTHOR, ("[TXTIn] WebVTT line %s: %s\n", line, message) );
	return e;
}

fixed #1897


Status Conv3DShape(shape_inference::InferenceContext* c) {
  ShapeHandle input_shape;
  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 5, &input_shape));
  ShapeHandle filter_shape;
  TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 5, &filter_shape));

  string data_format;
  Status s = c->GetAttr("data_format", &data_format);

  std::vector<int32> dilations;
  TF_RETURN_IF_ERROR(c->GetAttr("dilations", &dilations));

  if (dilations.size() != 5) {
    return errors::InvalidArgument(
        "Conv3D requires the dilation attribute to contain 5 values, but got: ",
        dilations.size());
  }

  std::vector<int32> strides;
  TF_RETURN_IF_ERROR(c->GetAttr("strides", &strides));
  if (strides.size() != 5) {
    return errors::InvalidArgument(
        "Conv3D requires the stride attribute to contain 5 values, but got: ",
        strides.size());
  }

  int32_t stride_planes, stride_rows, stride_cols;
  int32_t dilation_planes, dilation_rows, dilation_cols;
  if (s.ok() && data_format == "NCDHW") {
    // Convert input_shape to NDHWC.
    auto dim = [&](char dimension) {
      return c->Dim(input_shape, GetTensorDimIndex<3>(FORMAT_NCHW, dimension));
    };
    input_shape =
        c->MakeShape({{dim('N'), dim('0'), dim('1'), dim('2'), dim('C')}});
    stride_planes = strides[2];
    stride_rows = strides[3];
    stride_cols = strides[4];
    dilation_planes = dilations[2];
    dilation_cols = dilations[3];
    dilation_rows = dilations[4];
  } else {
    stride_planes = strides[1];
    stride_rows = strides[2];
    stride_cols = strides[3];
    dilation_planes = dilations[1];
    dilation_cols = dilations[2];
    dilation_rows = dilations[3];
  }

  DimensionHandle batch_size_dim = c->Dim(input_shape, 0);
  DimensionHandle in_planes_dim = c->Dim(input_shape, 1);
  DimensionHandle in_rows_dim = c->Dim(input_shape, 2);
  DimensionHandle in_cols_dim = c->Dim(input_shape, 3);
  DimensionHandle input_depth_dim = c->Dim(input_shape, 4);

  DimensionHandle filter_planes_dim = c->Dim(filter_shape, 0);
  DimensionHandle filter_rows_dim = c->Dim(filter_shape, 1);
  DimensionHandle filter_cols_dim = c->Dim(filter_shape, 2);
  DimensionHandle filter_input_depth_dim = c->Dim(filter_shape, 3);
  DimensionHandle output_depth_dim = c->Dim(filter_shape, 4);

  // Check that the input tensor and the filter tensor agree on the channel
  // count.
  if (c->ValueKnown(input_depth_dim) && c->ValueKnown(filter_input_depth_dim)) {
    int64_t input_depth_value = c->Value(input_depth_dim),
            filter_input_depth_value = c->Value(filter_input_depth_dim);
    if (filter_input_depth_value == 0)
      return errors::InvalidArgument("Depth of filter must not be 0");
    if (input_depth_value % filter_input_depth_value != 0)
      return errors::InvalidArgument(
          "Depth of input (", input_depth_value,
          ") is not a multiple of input depth of filter (",
          filter_input_depth_value, ")");
    if (input_depth_value != filter_input_depth_value) {
      int64_t num_groups = input_depth_value / filter_input_depth_value;
      if (c->ValueKnown(output_depth_dim)) {
        int64_t output_depth_value = c->Value(output_depth_dim);
        if (num_groups == 0)
          return errors::InvalidArgument("Number of groups must not be 0");
        if (output_depth_value % num_groups != 0)
          return errors::InvalidArgument(
              "Depth of output (", output_depth_value,
              ") is not a multiple of the number of groups (", num_groups, ")");
      }
    }
  }

  Padding padding;
  TF_RETURN_IF_ERROR(c->GetAttr("padding", &padding));
  DimensionHandle output_planes, output_rows, output_cols;

  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDimsV2(
      c, in_planes_dim, filter_planes_dim, dilation_planes, stride_planes,
      padding, -1, -1, &output_planes));
  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDimsV2(
      c, in_rows_dim, filter_rows_dim, dilation_rows, stride_rows, padding, -1,
      -1, &output_rows));
  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDimsV2(
      c, in_cols_dim, filter_cols_dim, dilation_cols, stride_cols, padding, -1,
      -1, &output_cols));

  ShapeHandle output_shape;
  if (data_format == "NCDHW") {
    output_shape = c->MakeShape({batch_size_dim, output_depth_dim,
                                 output_planes, output_rows, output_cols});
  } else {
    output_shape = c->MakeShape({batch_size_dim, output_planes, output_rows,
                                 output_cols, output_depth_dim});
  }
  c->set_output(0, output_shape);
  return Status::OK();
}

Fix Segfault in Concat V2 shape function.

PiperOrigin-RevId: 412120654
Change-Id: I3ff915faea694f9ad8b00024e9af2de9909011be


Status Conv2DShapeWithExplicitPadding(shape_inference::InferenceContext* c) {
  return Conv2DShapeImpl(c, true);
}


Status ValidateVariableResourceHandle(
    InferenceContext* c, std::vector<ShapeAndType>* shape_and_type) {
  auto* handle_data = c->input_handle_shapes_and_types(0);
  if (handle_data == nullptr || handle_data->empty()) {
    shape_and_type->emplace_back(c->UnknownShape(), DT_INVALID);
  } else {
    *shape_and_type = *handle_data;
    DataType value_dtype;
    TF_RETURN_IF_ERROR(c->GetAttr("dtype", &value_dtype));
    if (shape_and_type->at(0).dtype != value_dtype) {
      return errors::InvalidArgument(
          "Trying to read variable with wrong dtype. "
          "Expected ",
          DataTypeString(shape_and_type->at(0).dtype), " got ",
          DataTypeString(value_dtype));
    }
  }
  return Status::OK();
}


Status UnsortedSegmentReductionShapeFn(InferenceContext* c) {
  ShapeHandle s_data = c->input(0);
  ShapeHandle s_segment_ids = c->input(1);
  ShapeHandle s_num_segments = c->input(2);
  TF_RETURN_IF_ERROR(c->WithRank(s_num_segments, 0, &s_num_segments));

  ShapeHandle out;

  // Leading dimensions of data must be compatible with dimensions of
  // <s_segment_ids>.
  if (c->RankKnown(s_segment_ids)) {
    TF_RETURN_IF_ERROR(
        c->MergePrefix(s_data, s_segment_ids, &s_data, &s_segment_ids));

    // Get the value of the num_segments input tensor.
    DimensionHandle num_segments_dim;
    TF_RETURN_IF_ERROR(c->MakeDimForScalarInput(2, &num_segments_dim));

    // Output is {segment_id_rank} + s_data[segment_id_rank:].
    ShapeHandle s_data_suffix;
    TF_RETURN_IF_ERROR(
        c->Subshape(s_data, c->Rank(s_segment_ids), &s_data_suffix));
    TF_RETURN_IF_ERROR(
        c->Concatenate(c->Vector(num_segments_dim), s_data_suffix, &out));
  } else {
    out = c->UnknownShape();
  }
  c->set_output(0, out);
  return Status::OK();
}


Status GetWindowedOutputSizeFromDimsV2(
    shape_inference::InferenceContext* c,
    shape_inference::DimensionHandle input_size,
    shape_inference::DimensionOrConstant filter_size, int64_t dilation_rate,
    int64_t stride, Padding padding_type, int64_t padding_before,
    int64_t padding_after, shape_inference::DimensionHandle* output_size) {
  if (stride <= 0) {
    return errors::InvalidArgument("Stride must be > 0, but got ", stride);
  }

  if (dilation_rate < 1) {
    return errors::InvalidArgument("Dilation rate must be >= 1, but got ",
                                   dilation_rate);
  }

  // See also the parallel implementation in GetWindowedOutputSizeVerbose.
  switch (padding_type) {
    case Padding::VALID:
      padding_before = padding_after = 0;
      TF_FALLTHROUGH_INTENDED;
    case Padding::EXPLICIT:
      TF_RETURN_IF_ERROR(
          c->Add(input_size, padding_before + padding_after, &input_size));
      if (dilation_rate > 1) {
        DimensionHandle window_size;
        TF_RETURN_IF_ERROR(
            c->Subtract(c->MakeDim(filter_size), 1, &window_size));
        TF_RETURN_IF_ERROR(
            c->Multiply(window_size, dilation_rate, &window_size));
        TF_RETURN_IF_ERROR(c->Add(window_size, 1, &window_size));
        TF_RETURN_IF_ERROR(c->Subtract(input_size, window_size, output_size));
      } else {
        TF_RETURN_IF_ERROR(c->Subtract(input_size, filter_size, output_size));
      }
      TF_RETURN_IF_ERROR(c->Add(*output_size, stride, output_size));
      TF_RETURN_IF_ERROR(c->Divide(*output_size, stride,
                                   /*evenly_divisible=*/false, output_size));
      break;
    case Padding::SAME:
      TF_RETURN_IF_ERROR(c->Add(input_size, stride - 1, output_size));
      TF_RETURN_IF_ERROR(c->Divide(*output_size, stride,
                                   /*evenly_divisible=*/false, output_size));
      break;
  }
  return Status::OK();
}


Status ReductionShapeHelper(const Tensor* reduction_indices_t,
                            const int32_t input_rank,
                            std::set<int64_t>* true_indices) {
  auto reduction_indices = reduction_indices_t->flat<T>();
  for (int i = 0; i < reduction_indices_t->NumElements(); ++i) {
    const T reduction_index = reduction_indices(i);
    if (reduction_index < -input_rank || reduction_index >= input_rank) {
      return errors::InvalidArgument("Invalid reduction dimension ",
                                     reduction_index, " for input with ",
                                     input_rank, " dimensions.");
    }

    auto wrapped_index = reduction_index;
    if (wrapped_index < 0) {
      wrapped_index += input_rank;
    }

    true_indices->insert(wrapped_index);
  }
  return Status::OK();
}


Status GatherNdShape(InferenceContext* c) {
  ShapeHandle params;
  std::vector<ShapeAndType> handle_shape_and_type;
  if (c->input_handle_shapes_and_types(0) != nullptr) {
    TF_RETURN_IF_ERROR(
        ValidateVariableResourceHandle(c, &handle_shape_and_type));
    params = handle_shape_and_type[0].shape;
  } else {
    params = c->input(0);
  }
  ShapeHandle indices;
  TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(1), 1, &indices));
  DimensionHandle r_dim = c->Dim(indices, -1);

  if (!c->RankKnown(params) || !c->ValueKnown(r_dim)) {
    c->set_output(0, c->UnknownShape());
    return Status::OK();
  }

  if (c->Value(r_dim) > c->Rank(params)) {
    return errors::InvalidArgument(
        "indices.shape[-1] must be <= params.rank, but saw indices shape: ",
        c->DebugString(indices), " and params shape: ", c->DebugString(params));
  }

  // Remove r_dim from indices to get output.
  ShapeHandle indices_slice;
  ShapeHandle params_slice;
  TF_RETURN_IF_ERROR(c->Subshape(indices, 0, -1, &indices_slice));
  TF_RETURN_IF_ERROR(c->Subshape(params, c->Value(r_dim), &params_slice));
  ShapeHandle out;
  TF_RETURN_IF_ERROR(c->Concatenate(indices_slice, params_slice, &out));
  c->set_output(0, out);
  return Status::OK();
}


Status FusedBatchNormShape(shape_inference::InferenceContext* c) {
  string data_format_str;
  TF_RETURN_IF_ERROR(c->GetAttr("data_format", &data_format_str));
  TensorFormat data_format;
  if (!FormatFromString(data_format_str, &data_format)) {
    return errors::InvalidArgument("Invalid data format string: ",
                                   data_format_str);
  }
  const int rank =
      (data_format_str == "NDHWC" || data_format_str == "NCDHW") ? 5 : 4;
  ShapeHandle x;
  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), rank, &x));

  bool is_training;
  TF_RETURN_IF_ERROR(c->GetAttr("is_training", &is_training));
  float exponential_avg_factor;
  if (!c->GetAttr("exponential_avg_factor", &exponential_avg_factor).ok()) {
    exponential_avg_factor = 1.0f;  // default value
  }
  int number_inputs = (is_training && exponential_avg_factor == 1.0f) ? 3 : 5;

  int channel_dim_index = GetTensorFeatureDimIndex(rank, data_format);
  DimensionHandle channel_dim = c->Dim(x, channel_dim_index);

  // covers scale, offset, and if is_training is false, mean, variance
  for (int i = 1; i < number_inputs; ++i) {
    ShapeHandle vec;
    TF_RETURN_IF_ERROR(c->WithRank(c->input(i), 1, &vec));
    TF_RETURN_IF_ERROR(c->Merge(channel_dim, c->Dim(vec, 0), &channel_dim));
  }

  ShapeHandle y;
  TF_RETURN_IF_ERROR(c->ReplaceDim(x, channel_dim_index, channel_dim, &y));
  c->set_output(0, y);
  ShapeHandle vector_shape = c->Vector(channel_dim);
  c->set_output(1, vector_shape);
  c->set_output(2, vector_shape);
  c->set_output(3, vector_shape);
  c->set_output(4, vector_shape);
  return Status::OK();
}


Status MaxPoolV2Shape(shape_inference::InferenceContext* c, int num_inputs) {
  string data_format_str;
  TensorFormat data_format;
  Status s = c->GetAttr("data_format", &data_format_str);
  if (s.ok()) {
    FormatFromString(data_format_str, &data_format);
  } else {
    data_format = FORMAT_NHWC;
  }

  const int rank = (data_format == FORMAT_NCHW_VECT_C) ? 5 : 4;
  ShapeHandle input_shape;
  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), rank, &input_shape));

  TF_RETURN_IF_ERROR(
      CheckFormatConstraintsOnShape(data_format, input_shape, "input", c));

  std::vector<int32> kernel_sizes;
  std::vector<int32> strides;

  if (c->num_inputs() + 2 == num_inputs) {
    TF_RETURN_IF_ERROR(c->GetAttr("ksize", &kernel_sizes));

    TF_RETURN_IF_ERROR(c->GetAttr("strides", &strides));
  } else {
    // Verify shape of ksize and strides input.
    ShapeHandle size;
    DimensionHandle unused;
    TF_RETURN_IF_ERROR(c->WithRank(c->input(c->num_inputs() - 2), 1, &size));
    TF_RETURN_IF_ERROR(c->WithValue(c->Dim(size, 0), 4, &unused));
    TF_RETURN_IF_ERROR(c->WithRank(c->input(c->num_inputs() - 1), 1, &size));
    TF_RETURN_IF_ERROR(c->WithValue(c->Dim(size, 0), 4, &unused));

    const Tensor* kernel_sizes_tensor = c->input_tensor(c->num_inputs() - 2);
    if (kernel_sizes_tensor == nullptr) {
      c->set_output(0, c->UnknownShape());
      return Status::OK();
    }
    kernel_sizes.resize(kernel_sizes_tensor->shape().num_elements());
    auto kernel_sizes_vec = kernel_sizes_tensor->flat<int32>();
    std::copy_n(&kernel_sizes_vec(0), kernel_sizes.size(),
                kernel_sizes.begin());

    const Tensor* strides_tensor = c->input_tensor(c->num_inputs() - 1);
    if (strides_tensor == nullptr) {
      c->set_output(0, c->UnknownShape());
      return Status::OK();
    }
    strides.resize(strides_tensor->shape().num_elements());
    auto strides_vec = strides_tensor->flat<int32>();
    std::copy_n(&strides_vec(0), strides.size(), strides.begin());
  }

  if (strides.size() != 4) {
    return errors::InvalidArgument(
        "MaxPool requires the stride attribute to contain 4 values, but "
        "got: ",
        strides.size());
  }
  if (kernel_sizes.size() != 4) {
    return errors::InvalidArgument(
        "MaxPool requires the ksize attribute to contain 4 values, but got: ",
        kernel_sizes.size());
  }

  int32_t stride_depth = GetTensorDim(strides, data_format, 'C');
  int32_t stride_rows = GetTensorDim(strides, data_format, 'H');
  int32_t stride_cols = GetTensorDim(strides, data_format, 'W');
  int32_t kernel_depth = GetTensorDim(kernel_sizes, data_format, 'C');
  int32_t kernel_rows = GetTensorDim(kernel_sizes, data_format, 'H');
  int32_t kernel_cols = GetTensorDim(kernel_sizes, data_format, 'W');

  constexpr int num_spatial_dims = 2;
  DimensionHandle batch_size_dim = c->Dim(
      input_shape, GetTensorDimIndex<num_spatial_dims>(data_format, 'N'));
  DimensionHandle in_rows_dim = c->Dim(
      input_shape, GetTensorDimIndex<num_spatial_dims>(data_format, 'H'));
  DimensionHandle in_cols_dim = c->Dim(
      input_shape, GetTensorDimIndex<num_spatial_dims>(data_format, 'W'));
  DimensionHandle in_depth_dim = c->Dim(
      input_shape, GetTensorDimIndex<num_spatial_dims>(data_format, 'C'));

  Padding padding;
  TF_RETURN_IF_ERROR(c->GetAttr("padding", &padding));

  ShapeHandle output_shape;
  DimensionHandle output_rows, output_cols, output_depth;
  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDims(
      c, in_rows_dim, kernel_rows, stride_rows, padding, &output_rows));
  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDims(
      c, in_cols_dim, kernel_cols, stride_cols, padding, &output_cols));
  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDims(
      c, in_depth_dim, kernel_depth, stride_depth, padding, &output_depth));

  TF_RETURN_IF_ERROR(MakeShapeFromFormat(data_format, batch_size_dim,
                                         {output_rows, output_cols},
                                         output_depth, &output_shape, c));

  c->set_output(0, output_shape);
  return Status::OK();
}


Status Conv2DShape(shape_inference::InferenceContext* c) {
  return Conv2DShapeImpl(c, false);
}


Status SparseReduceShapeFn(InferenceContext* c) {
  // Input 0: input_indices
  // Input 1: input_values
  // Input 2: input_shape
  // Input 3: reduction_axes
  // Attr: keep_dims
  bool keep_dims = false;
  TF_RETURN_IF_ERROR(c->GetAttr("keep_dims", &keep_dims));

  const Tensor* shape_tensor = c->input_tensor(2);
  const Tensor* axes_tensor = c->input_tensor(3);
  if (shape_tensor != nullptr && axes_tensor != nullptr) {
    auto shape_vec = shape_tensor->flat<int64_t>();
    auto axes_vec = axes_tensor->flat<int32>();

    int64_t ndims = shape_vec.size();
    absl::flat_hash_set<int64_t> axes;
    if (ndims == 0)
      return errors::InvalidArgument(
          "Number of dims in shape tensor must not be 0");
    for (int i = 0; i < axes_vec.size(); i++) {
      axes.insert((axes_vec(i) + ndims) % ndims);
    }

    std::vector<DimensionHandle> dims;
    if (keep_dims) {
      dims.reserve(ndims);
      for (int d = 0; d < ndims; ++d) {
        if (axes.find(d) == axes.end()) {
          dims.push_back(c->MakeDim(shape_vec(d)));
        } else {
          dims.push_back(c->MakeDim(1));
        }
      }
    } else {
      for (int d = 0; d < ndims; ++d) {
        if (axes.find(d) == axes.end()) {
          dims.push_back(c->MakeDim(shape_vec(d)));
        }
      }
    }

    c->set_output(0, c->MakeShape(dims));
    return Status::OK();
  }
  return UnknownShape(c);
}


Status QuantizedConv2DShape(InferenceContext* c) {
  TF_RETURN_IF_ERROR(shape_inference::Conv2DShape(c));
  ShapeHandle unused;
  TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused));
  TF_RETURN_IF_ERROR(c->WithRank(c->input(3), 0, &unused));
  TF_RETURN_IF_ERROR(c->WithRank(c->input(4), 0, &unused));
  TF_RETURN_IF_ERROR(c->WithRank(c->input(5), 0, &unused));
  c->set_output(1, c->Scalar());
  c->set_output(2, c->Scalar());
  return Status::OK();
}


Status BiasAddShape(shape_inference::InferenceContext* c) {
  ShapeHandle input_shape;

  // Fetch the data_format attribute, which may not exist.
  string data_format;
  Status s = c->GetAttr("data_format", &data_format);

  if (s.ok() && data_format == "NCHW") {
    TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), 3, &input_shape));
  } else {
    TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), 2, &input_shape));
  }

  ShapeHandle bias_shape;
  TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 1, &bias_shape));
  DimensionHandle bias_dim = c->Dim(bias_shape, 0);

  // If rank unknown, return unknown shape.
  if (!c->RankKnown(input_shape)) {
    c->set_output(0, c->UnknownShape());
    return Status::OK();
  }

  // Output has the same shape as the input, and matches the length of
  // the bias in its bias dimension.
  ShapeHandle output_shape;
  if (s.ok() && data_format == "NCHW") {
    // Merge the length of bias_shape into the third to last dimension
    ShapeHandle first;
    TF_RETURN_IF_ERROR(c->Subshape(input_shape, 0, 1, &first));

    ShapeHandle last;
    TF_RETURN_IF_ERROR(c->Subshape(input_shape, 2, &last));

    DimensionHandle input_bias_dim = c->Dim(input_shape, 1);
    DimensionHandle merged_bias_dim;
    TF_RETURN_IF_ERROR(c->Merge(input_bias_dim, bias_dim, &merged_bias_dim));
    ShapeHandle merged_bias = c->Vector(merged_bias_dim);

    ShapeHandle temp;
    TF_RETURN_IF_ERROR(c->Concatenate(first, merged_bias, &temp));
    TF_RETURN_IF_ERROR(c->Concatenate(temp, last, &output_shape));
  } else {
    ShapeHandle all_but_bias;
    TF_RETURN_IF_ERROR(c->Subshape(input_shape, 0, -1, &all_but_bias));

    DimensionHandle input_bias_dim = c->Dim(input_shape, -1);
    DimensionHandle merged_bias_dim;
    TF_RETURN_IF_ERROR(c->Merge(input_bias_dim, bias_dim, &merged_bias_dim));

    ShapeHandle merged_bias = c->Vector(merged_bias_dim);
    TF_RETURN_IF_ERROR(
        c->Concatenate(all_but_bias, merged_bias, &output_shape));
  }

  c->set_output(0, output_shape);
  return Status::OK();
}


Status GetWindowedOutputSizeFromDims(
    shape_inference::InferenceContext* c,
    shape_inference::DimensionHandle input_size,
    shape_inference::DimensionOrConstant filter_size, int64_t stride,
    Padding padding_type, shape_inference::DimensionHandle* output_size) {
  if (padding_type == Padding::EXPLICIT) {
    return errors::Internal(
        "GetWindowedOutputSizeFromDims does not handle EXPLICIT padding; call "
        "GetWindowedOutputSizeFromDimsV2 instead");
  }
  return GetWindowedOutputSizeFromDimsV2(c, input_size, filter_size,
                                         /*dilation_rate=*/1, stride,
                                         padding_type,
                                         // Give dummy values of -1 to
                                         // padding_before and padding_after,
                                         // since explicit padding is not used.
                                         -1, -1, output_size);
}


Status BatchMatMulShape(shape_inference::InferenceContext* c) {
  ShapeHandle a_shape;
  ShapeHandle b_shape;
  TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), 2, &a_shape));
  TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(1), 2, &b_shape));

  // Determine output rows and cols.
  bool adj_x;
  bool adj_y;
  TF_RETURN_IF_ERROR(c->GetAttr("adj_x", &adj_x));
  TF_RETURN_IF_ERROR(c->GetAttr("adj_y", &adj_y));
  DimensionHandle output_rows = c->Dim(a_shape, adj_x ? -1 : -2);
  DimensionHandle output_cols = c->Dim(b_shape, adj_y ? -2 : -1);

  // Batch dims match between inputs.
  ShapeHandle a_batch_dims;
  ShapeHandle b_batch_dims;
  ShapeHandle batch_dims;
  TF_RETURN_IF_ERROR(c->Subshape(a_shape, 0, -2, &a_batch_dims));
  TF_RETURN_IF_ERROR(c->Subshape(b_shape, 0, -2, &b_batch_dims));
  TF_RETURN_IF_ERROR(c->Merge(a_batch_dims, b_batch_dims, &batch_dims));

  // Assert inner dims match.
  DimensionHandle unused;
  TF_RETURN_IF_ERROR(c->Merge(c->Dim(a_shape, adj_x ? -2 : -1),
                              c->Dim(b_shape, adj_y ? -1 : -2), &unused));

  ShapeHandle out;
  TF_RETURN_IF_ERROR(
      c->Concatenate(batch_dims, c->Matrix(output_rows, output_cols), &out));
  c->set_output(0, out);
  return Status::OK();
}


Status DatasetIteratorShape(shape_inference::InferenceContext* c) {
  shape_inference::ShapeHandle unused;
  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &unused));
  std::vector<PartialTensorShape> output_shapes;
  TF_RETURN_IF_ERROR(c->GetAttr("output_shapes", &output_shapes));
  const int output_shapes_size = output_shapes.size();
  if (output_shapes_size != c->num_outputs()) {
    return errors::InvalidArgument(
        "`output_shapes` must be the same length as `output_types` (",
        output_shapes.size(), " vs. ", c->num_outputs());
  }
  for (size_t i = 0; i < output_shapes.size(); ++i) {
    shape_inference::ShapeHandle output_shape_handle;
    TF_RETURN_IF_ERROR(c->MakeShapeFromPartialTensorShape(
        output_shapes[i], &output_shape_handle));
    c->set_output(static_cast<int>(i), output_shape_handle);
  }
  return Status::OK();
}


Status UnchangedShape(shape_inference::InferenceContext* c) {
  c->set_output(0, c->input(0));
  auto* handle_data = c->input_handle_shapes_and_types(0);
  if (handle_data != nullptr) {
    c->set_output_handle_shapes_and_types(0, *handle_data);
  }
  return Status::OK();
}


Status Conv2DBackpropFilterWithBiasShape(shape_inference::InferenceContext* c) {
  ShapeHandle input_shape;
  // Fetch the data_format attribute, which may not exist.
  string data_format;
  Status s = c->GetAttr("data_format", &data_format);

  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 4, &input_shape));
  if (s.ok() && data_format == "NCHW") {
    c->set_output(1, c->Vector(c->Dim(input_shape, -3)));
  } else {
    c->set_output(1, c->Vector(c->Dim(input_shape, -1)));
  }
  ShapeHandle sh;
  TF_RETURN_IF_ERROR(c->MakeShapeFromShapeTensor(1, &sh));
  TF_RETURN_IF_ERROR(c->WithRank(sh, 4, &sh));
  c->set_output(0, sh);
  return Status::OK();
}


Status BatchMatMulV2Shape(shape_inference::InferenceContext* c) {
  ShapeHandle a_shape;
  ShapeHandle b_shape;
  TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), 2, &a_shape));
  TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(1), 2, &b_shape));

  // Determine output rows and columns.
  bool adj_x;
  bool adj_y;
  TF_RETURN_IF_ERROR(c->GetAttr("adj_x", &adj_x));
  TF_RETURN_IF_ERROR(c->GetAttr("adj_y", &adj_y));
  DimensionHandle output_rows = c->Dim(a_shape, adj_x ? -1 : -2);
  DimensionHandle output_cols = c->Dim(b_shape, adj_y ? -2 : -1);

  // Inner dimensions should be compatible.
  DimensionHandle inner_merged;
  TF_RETURN_IF_ERROR(c->Merge(c->Dim(a_shape, adj_x ? -2 : -1),
                              c->Dim(b_shape, adj_y ? -1 : -2), &inner_merged));

  // Batch dimensions should broadcast with each other.
  ShapeHandle a_batch_shape;
  ShapeHandle b_batch_shape;
  ShapeHandle output_batch_shape;
  TF_RETURN_IF_ERROR(c->Subshape(a_shape, 0, -2, &a_batch_shape));
  TF_RETURN_IF_ERROR(c->Subshape(b_shape, 0, -2, &b_batch_shape));

  TF_RETURN_IF_ERROR(BroadcastBinaryOpOutputShapeFnHelper(
      c, a_batch_shape, b_batch_shape, true, &output_batch_shape));

  ShapeHandle output_shape;
  TF_RETURN_IF_ERROR(c->Concatenate(
      output_batch_shape, c->Matrix(output_rows, output_cols), &output_shape));

  c->set_output(0, output_shape);
  return Status::OK();
}


Status MakeShapeFromFormat(TensorFormat format, DimensionOrConstant N,
                           const std::vector<DimensionOrConstant>& spatial,
                           DimensionOrConstant C, ShapeHandle* out,
                           shape_inference::InferenceContext* context) {
  const int num_dims = GetTensorDimsFromSpatialDims(spatial.size(), format);
  std::vector<DimensionHandle> dims_actual(num_dims);
  dims_actual[GetTensorBatchDimIndex(num_dims, format)] = context->MakeDim(N);
  int outer_c_index = GetTensorFeatureDimIndex(num_dims, format);
  dims_actual[outer_c_index] = context->MakeDim(C);
  if (format == FORMAT_NCHW_VECT_C) {
    dims_actual[GetTensorInnerFeatureDimIndex(num_dims, format)] =
        context->MakeDim(4);
  } else if (format == FORMAT_NHWC_VECT_W) {
    dims_actual[GetTensorInnerWidthDimIndex(num_dims, format)] =
        context->MakeDim(4);
  }
  for (int spatial_dim = 0, end = spatial.size(); spatial_dim < end;
       spatial_dim++) {
    dims_actual[GetTensorSpatialDimIndex(num_dims, format, spatial_dim)] =
        context->MakeDim(spatial[spatial_dim]);
  }
  *out = context->MakeShape(dims_actual);
  return Status::OK();
}


Status SliceHelper(InferenceContext* c, ShapeHandle begin_value,
                   const Tensor* sizes_value,
                   std::vector<DimensionHandle>* dims) {
  auto sizes_vec = sizes_value->vec<T>();
  for (int i = 0; i < sizes_value->NumElements(); ++i) {
    DimensionHandle dim = c->Dim(c->input(0), i);
    if (sizes_vec(i) != -1) {
      auto dim_val = c->Value(dim);
      if (sizes_vec(i) < 0) {
        return errors::InvalidArgument(
            "Out of bounds slicing on dimension ", i, " of length ", dim_val,
            ": sizes vector cannot be < -1, but was ", sizes_vec(i));
      }

      dims->emplace_back(c->MakeDim(sizes_vec(i)));
    } else {
      DimensionHandle result;
      TF_RETURN_IF_ERROR(c->Subtract(dim, c->Dim(begin_value, i), &result));
      dims->emplace_back(result);
    }
  }

  return Status::OK();
}


Status MaxPoolGradShape(shape_inference::InferenceContext* c) {
  return UnchangedShapeWithRank(c, 4);
}


Status QuantizedConcatV2Shape(InferenceContext* c, int num_inputs_to_concat) {
  return ConcatShapeHelper(c, 0 /* start_value_index */,
                           num_inputs_to_concat /* end_value_index */,
                           num_inputs_to_concat /* dim_index */);
}


Status ShapeFromDimensions(DimensionHandle batch_dim,
                           gtl::ArraySlice<DimensionHandle> spatial_dims,
                           DimensionHandle filter_dim, TensorFormat format,
                           absl::optional<DimensionHandle> vect_size,
                           InferenceContext* context, ShapeHandle* shape) {
  const int32_t rank =
      GetTensorDimsFromSpatialDims(spatial_dims.size(), format);
  std::vector<DimensionHandle> out_dims(rank);

  // Batch.
  out_dims[tensorflow::GetTensorBatchDimIndex(rank, format)] = batch_dim;
  // Spatial.
  for (int spatial_dim_index = 0, end = spatial_dims.size();
       spatial_dim_index < end; ++spatial_dim_index) {
    out_dims[tensorflow::GetTensorSpatialDimIndex(
        rank, format, spatial_dim_index)] = spatial_dims[spatial_dim_index];
  }
  // Channel.
  if (format == tensorflow::FORMAT_NCHW_VECT_C) {
    // When format is NCHW_VECT_C, factor the feature map count into the outer
    // feature count and the inner feature count (4 or 32).
    CHECK(vect_size.has_value());  // Crash ok.
    TF_RETURN_IF_ERROR(context->Divide(
        filter_dim, *vect_size, /*evenly_divisible=*/true,
        &out_dims[tensorflow::GetTensorFeatureDimIndex(rank, format)]));
    out_dims[GetTensorInnerFeatureDimIndex(rank, format)] = *vect_size;
  } else {
    out_dims[tensorflow::GetTensorFeatureDimIndex(rank, format)] = filter_dim;
  }

  *shape = context->MakeShape(out_dims);
  return tensorflow::Status::OK();
}


Status ExplicitShapes(InferenceContext* c) {
  std::vector<PartialTensorShape> shapes;
  TF_RETURN_IF_ERROR(c->GetAttr("shapes", &shapes));
  if (shapes.empty()) {
    return errors::Internal("shapes attribute is empty");
  }
  for (int i = 0, end = shapes.size(); i < end; ++i) {
    ShapeHandle output_shape;
    TF_RETURN_IF_ERROR(
        c->MakeShapeFromPartialTensorShape(shapes[i], &output_shape));
    c->set_output(i, output_shape);
  }
  return Status::OK();
}


Status MatrixSetDiagV2Shape(shape_inference::InferenceContext* c) {
  ShapeHandle input_shape, diag_shape, diag_index_shape;
  TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), 2, &input_shape));
  TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(1), 1, &diag_shape));
  TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(2), 1, &diag_index_shape));

  int32_t lower_diag_index = 0;
  int32_t upper_diag_index = 0;
  bool diag_index_known = false;
  const Tensor* diag_index_tensor = c->input_tensor(2);
  if (diag_index_tensor != nullptr && c->FullyDefined(diag_index_shape)) {
    diag_index_known = true;
    TF_RETURN_IF_ERROR(ReadDiagIndex(c, diag_index_tensor, &lower_diag_index,
                                     &upper_diag_index));
    if (lower_diag_index > upper_diag_index) {
      return errors::InvalidArgument(
          "lower_diag_index is greater than upper_diag_index");
    }
  }

  // Do more checks when input rank is known.
  if (c->RankKnown(input_shape)) {
    int32_t input_rank = c->Rank(input_shape);

    // If diag_index is set, we know the exact rank of diagonal.
    if (diag_index_known) {
      TF_RETURN_IF_ERROR(c->WithRank(
          c->input(1),
          (lower_diag_index == upper_diag_index) ? input_rank - 1 : input_rank,
          &diag_shape));
    } else {
      TF_RETURN_IF_ERROR(
          c->WithRankAtLeast(c->input(1), input_rank - 1, &diag_shape));
      TF_RETURN_IF_ERROR(
          c->WithRankAtMost(c->input(1), input_rank, &diag_shape));
    }

    // Validates lower_diag_index and upper_diag_index.
    const int32_t num_rows = c->Value(c->Dim(input_shape, input_rank - 2));
    const int32_t num_cols = c->Value(c->Dim(input_shape, input_rank - 1));
    if (num_rows != InferenceContext::kUnknownDim &&
        num_cols != InferenceContext::kUnknownDim) {
      if (lower_diag_index != 0 &&  // For when num_rows or num_cols == 0.
          (-num_rows >= lower_diag_index || lower_diag_index >= num_cols)) {
        return errors::InvalidArgument("lower_diag_index is out of bound.");
      }
      if (upper_diag_index != 0 &&  // For when num_rows or num_cols == 0.
          (-num_rows >= upper_diag_index || upper_diag_index >= num_cols)) {
        return errors::InvalidArgument("upper_diag_index is out of bound.");
      }
    }
  }

  ShapeHandle output_shape = input_shape;
  if (c->RankKnown(diag_shape) && !c->FullyDefined(input_shape)) {
    // Try to infer parts of shape from diag.
    ShapeHandle diag_prefix;
    TF_RETURN_IF_ERROR(c->Subshape(
        diag_shape, 0, (lower_diag_index == upper_diag_index) ? -1 : -2,
        &diag_prefix));

    // The inner matrices can be rectangular, so we can't pinpoint their
    // exact height and width by just lower_diag_index, upper_diag_index,
    // and the longest length of given diagonals.
    TF_RETURN_IF_ERROR(
        c->Concatenate(diag_prefix, c->UnknownShapeOfRank(2), &diag_shape));
    TF_RETURN_IF_ERROR(c->Merge(input_shape, diag_shape, &output_shape));
  }
  c->set_output(0, output_shape);
  return Status::OK();
}


Status DepthwiseConv2DNativeShape(shape_inference::InferenceContext* c) {
  return DepthwiseConv2DNativeShapeImpl(c, false);
}


Status FusedBatchNormGradExShape(shape_inference::InferenceContext* c) {
  TF_RETURN_IF_ERROR(FusedBatchNormGradShape(c));

  int num_side_inputs;
  TF_RETURN_IF_ERROR(c->GetAttr("num_side_inputs", &num_side_inputs));
  if (num_side_inputs == 0) {
    return Status::OK();
  }

  string data_format_str;
  TF_RETURN_IF_ERROR(c->GetAttr("data_format", &data_format_str));
  TensorFormat data_format;
  if (!FormatFromString(data_format_str, &data_format)) {
    return errors::InvalidArgument("Invalid data format string: ",
                                   data_format_str);
  }
  const int rank =
      (data_format_str == "NDHWC" || data_format_str == "NCDHW") ? 5 : 4;
  ShapeHandle y_backprop;
  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), rank, &y_backprop));
  ShapeHandle x;
  TF_RETURN_IF_ERROR(c->WithRank(c->input(1), rank, &x));

  int channel_dim_index = GetTensorFeatureDimIndex(rank, data_format);
  DimensionHandle channel_dim = c->Dim(y_backprop, channel_dim_index);
  TF_RETURN_IF_ERROR(
      c->Merge(channel_dim, c->Dim(x, channel_dim_index), &channel_dim));

  ShapeHandle side_input_backprop;
  TF_RETURN_IF_ERROR(c->ReplaceDim(y_backprop, channel_dim_index, channel_dim,
                                   &side_input_backprop));

  c->set_output(5, side_input_backprop);
  return Status::OK();
}


Status AvgPool3DGradShape(shape_inference::InferenceContext* c) {
  ShapeHandle s;
  TF_RETURN_IF_ERROR(c->MakeShapeFromShapeTensor(0, &s));
  TF_RETURN_IF_ERROR(c->WithRank(s, 5, &s));
  c->set_output(0, s);
  return Status::OK();
}


Status ScatterNdShapeHelper(InferenceContext* c, ShapeHandle indices_shape,
                            ShapeHandle updates_shape,
                            ShapeHandle input_shape) {
  if (c->Value(c->NumElements(input_shape)) == 0 &&
      (c->Value(c->NumElements(indices_shape)) > 0 ||
       c->Value(c->NumElements(updates_shape)) > 0)) {
    return errors::InvalidArgument(
        "Indices and updates specified for empty input");
  }

  if (c->RankKnown(indices_shape) && c->RankKnown(updates_shape)) {
    const int64_t outer_dims = c->Rank(indices_shape) - 1;
    const DimensionHandle ixdim = c->Dim(indices_shape, -1);

    // We can only do more validation if the last dimension of indices
    // is a known value.
    if (c->ValueKnown(ixdim)) {
      int64_t ix = c->Value(ixdim);
      ShapeHandle unused;
      ShapeHandle prefix_indices;
      TF_RETURN_IF_ERROR(
          c->Subshape(indices_shape, 0, outer_dims, &prefix_indices));
      ShapeHandle prefix_updates;
      TF_RETURN_IF_ERROR(
          c->Subshape(updates_shape, 0, outer_dims, &prefix_updates));

      Status s = c->Merge(prefix_indices, prefix_updates, &unused);
      if (!s.ok()) {
        return errors::InvalidArgument(
            "Dimensions [0,", outer_dims,
            ") of indices[shape=", c->DebugString(indices_shape),
            "] = ", c->DebugString(prefix_indices),
            " must match dimensions [0,", outer_dims,
            ") of updates[shape=", c->DebugString(updates_shape),
            "] = ", c->DebugString(prefix_updates), ": ", s.error_message());
      }

      ShapeHandle suffix_output;
      TF_RETURN_IF_ERROR(c->Subshape(input_shape, ix, &suffix_output));
      ShapeHandle suffix_updates;
      TF_RETURN_IF_ERROR(
          c->Subshape(updates_shape, outer_dims, &suffix_updates));
      s = c->Merge(suffix_output, suffix_updates, &unused);
      if (!s.ok()) {
        return errors::InvalidArgument(
            "Dimensions [", ix, ",", c->Rank(input_shape),
            ") of input[shape=", c->DebugString(input_shape),
            "] = ", c->DebugString(suffix_output), " must match dimensions [",
            outer_dims, ",", c->Rank(updates_shape),
            ") of updates[shape=", c->DebugString(updates_shape),
            "] = ", c->DebugString(suffix_updates), ": ", s.error_message());
      }
    }
  }

  if (c->input_handle_shapes_and_types(0) == nullptr && c->num_outputs() > 0) {
    // This is called for tf.scatter_nd; output is a tensor with this shape.
    c->set_output(0, input_shape);
  }
  return Status::OK();
}


Status DepthwiseConv2DNativeShapeWithExplicitPadding(
    shape_inference::InferenceContext* c) {
  return DepthwiseConv2DNativeShapeImpl(c, true);
}


Status ReadDiagIndex(InferenceContext* c, const Tensor* diag_index_tensor,
                     int32* lower_diag_index, int32* upper_diag_index) {
  // This function assumes that the shape of diag_index_tensor is fully defined.
  if (diag_index_tensor->dims() == 0) {
    *lower_diag_index = diag_index_tensor->scalar<int32>()();
    *upper_diag_index = *lower_diag_index;
  } else {
    int32_t num_elements = diag_index_tensor->dim_size(0);
    if (num_elements == 1) {
      *lower_diag_index = diag_index_tensor->vec<int32>()(0);
      *upper_diag_index = *lower_diag_index;
    } else if (num_elements == 2) {
      *lower_diag_index = diag_index_tensor->vec<int32>()(0);
      *upper_diag_index = diag_index_tensor->vec<int32>()(1);
    } else {
      return errors::InvalidArgument(
          "diag_index must be a vector with one or two elements. It has ",
          num_elements, " elements.");
    }
  }
  return Status::OK();
}


Status MaxPoolShapeImpl(shape_inference::InferenceContext* c,
                        bool supports_explicit_padding) {
  string data_format_str;
  TensorFormat data_format;
  Status s = c->GetAttr("data_format", &data_format_str);
  if (s.ok()) {
    FormatFromString(data_format_str, &data_format);
  } else {
    data_format = FORMAT_NHWC;
  }

  const int rank = (data_format == FORMAT_NCHW_VECT_C) ? 5 : 4;
  ShapeHandle input_shape;
  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), rank, &input_shape));

  TF_RETURN_IF_ERROR(
      CheckFormatConstraintsOnShape(data_format, input_shape, "input", c));

  std::vector<int32> strides;
  TF_RETURN_IF_ERROR(c->GetAttr("strides", &strides));
  if (strides.size() != 4) {
    return errors::InvalidArgument(
        "MaxPool requires the stride attribute to contain 4 values, but got: ",
        strides.size());
  }

  std::vector<int32> kernel_sizes;
  TF_RETURN_IF_ERROR(c->GetAttr("ksize", &kernel_sizes));
  if (kernel_sizes.size() != 4) {
    return errors::InvalidArgument(
        "MaxPool requires the ksize attribute to contain 4 values, but got: ",
        kernel_sizes.size());
  }

  int32_t stride_depth = GetTensorDim(strides, data_format, 'C');
  int32_t stride_rows = GetTensorDim(strides, data_format, 'H');
  int32_t stride_cols = GetTensorDim(strides, data_format, 'W');
  int32_t kernel_depth = GetTensorDim(kernel_sizes, data_format, 'C');
  int32_t kernel_rows = GetTensorDim(kernel_sizes, data_format, 'H');
  int32_t kernel_cols = GetTensorDim(kernel_sizes, data_format, 'W');

  constexpr int num_spatial_dims = 2;
  DimensionHandle batch_size_dim = c->Dim(
      input_shape, GetTensorDimIndex<num_spatial_dims>(data_format, 'N'));
  DimensionHandle in_rows_dim = c->Dim(
      input_shape, GetTensorDimIndex<num_spatial_dims>(data_format, 'H'));
  DimensionHandle in_cols_dim = c->Dim(
      input_shape, GetTensorDimIndex<num_spatial_dims>(data_format, 'W'));
  DimensionHandle in_depth_dim = c->Dim(
      input_shape, GetTensorDimIndex<num_spatial_dims>(data_format, 'C'));

  Padding padding;
  TF_RETURN_IF_ERROR(c->GetAttr("padding", &padding));

  std::vector<int64_t> explicit_paddings;
  if (supports_explicit_padding) {
    Status status = c->GetAttr("explicit_paddings", &explicit_paddings);
    // Use the default value, which is an empty list, if the attribute is not
    // found. Otherwise return the error to the caller.
    if (!status.ok() && !errors::IsNotFound(status)) {
      return status;
    }
    TF_RETURN_IF_ERROR(CheckValidPadding(padding, explicit_paddings,
                                         /*num_dims=*/4, data_format));
  } else {
    DCHECK(padding != Padding::EXPLICIT);
  }

  ShapeHandle output_shape;
  DimensionHandle output_rows, output_cols, output_depth;
  int64_t pad_rows_before = -1, pad_rows_after = -1;
  int64_t pad_cols_before = -1, pad_cols_after = -1;
  if (padding == Padding::EXPLICIT) {
    GetExplicitPaddingForDim(explicit_paddings, data_format, 'H',
                             &pad_rows_before, &pad_rows_after);
    GetExplicitPaddingForDim(explicit_paddings, data_format, 'W',
                             &pad_cols_before, &pad_cols_after);
  }
  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDimsV2(
      c, in_rows_dim, kernel_rows, /*dilation_rate=*/1, stride_rows, padding,
      pad_rows_before, pad_rows_after, &output_rows));
  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDimsV2(
      c, in_cols_dim, kernel_cols, /*dilation_rate=*/1, stride_cols, padding,
      pad_cols_before, pad_cols_after, &output_cols));
  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDimsV2(
      c, in_depth_dim, kernel_depth, /*dilation_rate=*/1, stride_depth, padding,
      /*pad_before*/ 0, /*pad_after*/ 0, &output_depth));

  TF_RETURN_IF_ERROR(MakeShapeFromFormat(data_format, batch_size_dim,
                                         {output_rows, output_cols},
                                         output_depth, &output_shape, c));

  c->set_output(0, output_shape);
  return Status::OK();
}


Status Pool3DShape(shape_inference::InferenceContext* c) {
  ShapeHandle input_shape;
  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 5, &input_shape));

  string data_format;
  Status s = c->GetAttr("data_format", &data_format);

  std::vector<int32> strides;
  TF_RETURN_IF_ERROR(c->GetAttr("strides", &strides));
  if (strides.size() != 5) {
    return errors::InvalidArgument(
        "Pool3D ops require the stride attribute to contain 5 values, but "
        "got: ",
        strides.size());
  }

  std::vector<int32> kernel_sizes;
  TF_RETURN_IF_ERROR(c->GetAttr("ksize", &kernel_sizes));
  if (kernel_sizes.size() != 5) {
    return errors::InvalidArgument(
        "Pool3D requires the ksize attribute to contain 5 values, but got: ",
        kernel_sizes.size());
  }

  int32_t stride_planes, stride_rows, stride_cols;
  int32_t kernel_planes, kernel_rows, kernel_cols;

  if (s.ok() && data_format == "NCDHW") {
    // Convert input_shape to NDHWC.
    auto dim = [&](char dimension) {
      return c->Dim(input_shape, GetTensorDimIndex<3>(FORMAT_NCHW, dimension));
    };
    input_shape =
        c->MakeShape({{dim('N'), dim('0'), dim('1'), dim('2'), dim('C')}});
    stride_planes = strides[2];
    stride_rows = strides[3];
    stride_cols = strides[4];
    kernel_planes = kernel_sizes[2];
    kernel_rows = kernel_sizes[3];
    kernel_cols = kernel_sizes[4];
  } else {
    stride_planes = strides[1];
    stride_rows = strides[2];
    stride_cols = strides[3];
    kernel_planes = kernel_sizes[1];
    kernel_rows = kernel_sizes[2];
    kernel_cols = kernel_sizes[3];
  }

  DimensionHandle batch_size_dim = c->Dim(input_shape, 0);
  DimensionHandle in_planes_dim = c->Dim(input_shape, 1);
  DimensionHandle in_rows_dim = c->Dim(input_shape, 2);
  DimensionHandle in_cols_dim = c->Dim(input_shape, 3);
  DimensionHandle output_depth_dim = c->Dim(input_shape, 4);

  Padding padding;
  TF_RETURN_IF_ERROR(c->GetAttr("padding", &padding));

  // TODO(mrry,shlens): Raise an error if the stride would cause
  // information in the input to be ignored. This will require a change
  // in the kernel implementation.
  DimensionHandle output_planes, output_rows, output_cols;
  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDims(
      c, in_planes_dim, kernel_planes, stride_planes, padding, &output_planes));
  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDims(
      c, in_rows_dim, kernel_rows, stride_rows, padding, &output_rows));
  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDims(
      c, in_cols_dim, kernel_cols, stride_cols, padding, &output_cols));

  ShapeHandle output_shape;
  if (data_format == "NCDHW") {
    output_shape = c->MakeShape({batch_size_dim, output_depth_dim,
                                 output_planes, output_rows, output_cols});
  } else {
    output_shape = c->MakeShape({batch_size_dim, output_planes, output_rows,
                                 output_cols, output_depth_dim});
  }

  c->set_output(0, output_shape);
  return Status::OK();
}


Status DimensionsFromShape(ShapeHandle shape, TensorFormat format,
                           DimensionHandle* batch_dim,
                           gtl::MutableArraySlice<DimensionHandle> spatial_dims,
                           DimensionHandle* filter_dim,
                           InferenceContext* context) {
  const int32_t rank =
      GetTensorDimsFromSpatialDims(spatial_dims.size(), format);
  // Batch.
  *batch_dim = context->Dim(shape, GetTensorBatchDimIndex(rank, format));
  // Spatial.
  for (int spatial_dim_index = 0, end = spatial_dims.size();
       spatial_dim_index < end; ++spatial_dim_index) {
    spatial_dims[spatial_dim_index] = context->Dim(
        shape, GetTensorSpatialDimIndex(rank, format, spatial_dim_index));
  }
  // Channel.
  *filter_dim = context->Dim(shape, GetTensorFeatureDimIndex(rank, format));
  if (format == FORMAT_NCHW_VECT_C) {
    TF_RETURN_IF_ERROR(context->Multiply(
        *filter_dim,
        context->Dim(shape, GetTensorInnerFeatureDimIndex(rank, format)),
        filter_dim));
  }
  return Status::OK();
}


Status MatMulShape(shape_inference::InferenceContext* c) {
  ShapeHandle a;
  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 2, &a));

  ShapeHandle b;
  TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 2, &b));

  bool transpose_a, transpose_b;
  TF_RETURN_IF_ERROR(c->GetAttr("transpose_a", &transpose_a));
  TF_RETURN_IF_ERROR(c->GetAttr("transpose_b", &transpose_b));
  DimensionHandle output_rows = transpose_a ? c->Dim(a, 1) : c->Dim(a, 0);
  DimensionHandle output_cols = transpose_b ? c->Dim(b, 0) : c->Dim(b, 1);

  // Validate that the inner shapes are compatible.
  DimensionHandle inner_a = transpose_a ? c->Dim(a, 0) : c->Dim(a, 1);
  DimensionHandle inner_b = transpose_b ? c->Dim(b, 1) : c->Dim(b, 0);
  DimensionHandle merged;
  TF_RETURN_IF_ERROR(c->Merge(inner_a, inner_b, &merged));

  c->set_output(0, c->Matrix(output_rows, output_cols));
  return Status::OK();
}


Status FusedBatchNormV3Shape(shape_inference::InferenceContext* c) {
  TF_RETURN_IF_ERROR(FusedBatchNormShape(c));
  c->set_output(5, c->UnknownShape());
  return Status::OK();
}


Status MaxPoolShapeWithExplicitPadding(shape_inference::InferenceContext* c) {
  return MaxPoolShapeImpl(c, /*supports_explicit_padding=*/true);
}


Status CheckFormatConstraintsOnShape(const TensorFormat tensor_format,
                                     const ShapeHandle shape_handle,
                                     const string& tensor_name,
                                     shape_inference::InferenceContext* c) {
  if (tensor_format == FORMAT_NCHW_VECT_C) {
    // Check that the vect dim has size 4 or 32.
    const int num_dims = c->Rank(shape_handle);
    DimensionHandle vect_dim = c->Dim(
        shape_handle, GetTensorInnerFeatureDimIndex(num_dims, tensor_format));
    int64_t vect_dim_val = c->Value(vect_dim);
    if (vect_dim_val != 4 && vect_dim_val != 32) {
      return errors::InvalidArgument(
          "VECT_C dimension must be 4 or 32, but is ", vect_dim_val);
    }
  }

  return Status::OK();
}


Status ExplicitShape(InferenceContext* c) {
  PartialTensorShape shape;
  TF_RETURN_IF_ERROR(c->GetAttr("shape", &shape));
  ShapeHandle output_shape;
  TF_RETURN_IF_ERROR(c->MakeShapeFromPartialTensorShape(shape, &output_shape));
  c->set_output(0, output_shape);
  return Status::OK();
}


Status AvgPoolGradShape(shape_inference::InferenceContext* c) {
  ShapeHandle s;
  TF_RETURN_IF_ERROR(c->MakeShapeFromShapeTensor(0, &s));
  TF_RETURN_IF_ERROR(c->WithRank(s, 4, &s));
  c->set_output(0, s);
  return Status::OK();
}


Status UnknownShape(shape_inference::InferenceContext* c) {
  for (int i = 0; i < c->num_outputs(); ++i) {
    c->set_output(i, c->UnknownShape());
  }
  return Status::OK();
}


Status MatrixDiagV2Shape(shape_inference::InferenceContext* c) {
  // Checks input ranks.
  ShapeHandle input_shape, diag_index_shape, unused_shape;
  TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), 1, &input_shape));
  TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(1), 1, &diag_index_shape));
  TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused_shape));
  TF_RETURN_IF_ERROR(c->WithRank(c->input(3), 0, &unused_shape));
  TF_RETURN_IF_ERROR(c->WithRank(c->input(4), 0, &unused_shape));

  // Reads the diagonal indices.
  const Tensor* diag_index_tensor = c->input_tensor(1);
  if (!c->RankKnown(input_shape) || !c->FullyDefined(diag_index_shape) ||
      diag_index_tensor == nullptr) {
    c->set_output(0, c->UnknownShape());
    return Status::OK();
  }
  int32_t lower_diag_index = 0;
  int32_t upper_diag_index = 0;
  TF_RETURN_IF_ERROR(ReadDiagIndex(c, diag_index_tensor, &lower_diag_index,
                                   &upper_diag_index));
  if (lower_diag_index > upper_diag_index) {
    return errors::InvalidArgument(
        "lower_diag_index is greater than upper_diag_index");
  }

  // Checks if the number of diagonals provided matches what we imply from
  // lower_diag_index and upper_diag_index.
  const int32_t input_rank = c->Rank(input_shape);
  if (lower_diag_index < upper_diag_index) {
    const int32_t num_diags = c->Value(c->Dim(input_shape, input_rank - 2));
    const int32_t other_dim = c->Value(c->Dim(input_shape, input_rank - 1));

    if (num_diags != (upper_diag_index - lower_diag_index + 1)) {
      return errors::InvalidArgument(
          "The number of rows of `diagonal` doesn't match the number of "
          "diagonals implied from `d_lower` and `d_upper`.\n",
          "num_diags = ", num_diags, ", d_lower = ", lower_diag_index,
          ", d_upper = ", upper_diag_index, " ", input_rank, " ", other_dim);
    }
  }

  // Reads num_rows and num_cols.
  const Tensor* num_rows_tensor = c->input_tensor(2);
  const Tensor* num_cols_tensor = c->input_tensor(3);
  int64_t num_rows = -1;
  int64_t num_cols = -1;
  if (num_rows_tensor != nullptr) {
    TF_RETURN_IF_ERROR(c->GetScalarFromTensor(num_rows_tensor, &num_rows));
  }
  if (num_cols_tensor != nullptr) {
    TF_RETURN_IF_ERROR(c->GetScalarFromTensor(num_cols_tensor, &num_cols));
  }

  // Infers the missing num_rows or num_cols: If both are missing, assume
  // output is square. Otherwise, use the smallest possible value. Also
  // validates the provided values.
  const int32_t max_diag_len = c->Value(c->Dim(input_shape, input_rank - 1));
  const int32_t min_num_rows = max_diag_len - std::min(upper_diag_index, 0);
  const int32_t min_num_cols = max_diag_len + std::max(lower_diag_index, 0);
  if (num_rows == -1 && num_cols == -1) {  // Special case.
    num_rows = std::max(min_num_rows, min_num_cols);
    num_cols = num_rows;
  }
  if (num_rows == -1) {
    num_rows = min_num_rows;
  } else if (num_rows < min_num_rows) {
    return errors::InvalidArgument("num_rows is too small");
  }
  if (num_cols == -1) {
    num_cols = min_num_cols;
  } else if (num_cols < min_num_cols) {
    return errors::InvalidArgument("num_cols is too small.");
  }
  // At least one of them must match the minimum length.
  if (num_rows != min_num_rows && num_cols != min_num_cols) {
    return errors::InvalidArgument(
        "num_rows and num_cols are not consistent with lower_diag_index, "
        "upper_diag_index, and the length of the given diagonals.\n",
        "num_rows = ", num_rows, " != min_num_rows = ", min_num_rows,
        ", num_cols = ", num_cols, " != min_num_cols = ", min_num_cols);
  }

  // Sets output shape.
  ShapeHandle output_shape;
  const DimensionHandle output_row_dim = c->MakeDim(num_rows);
  const DimensionHandle output_col_dim = c->MakeDim(num_cols);
  if (lower_diag_index == upper_diag_index) {
    TF_RETURN_IF_ERROR(c->ReplaceDim(input_shape, input_rank - 1,
                                     output_row_dim, &output_shape));
    TF_RETURN_IF_ERROR(
        c->Concatenate(output_shape, c->Vector(output_col_dim), &output_shape));
  } else {
    TF_RETURN_IF_ERROR(c->ReplaceDim(input_shape, input_rank - 2,
                                     output_row_dim, &output_shape));
    TF_RETURN_IF_ERROR(c->ReplaceDim(output_shape, input_rank - 1,
                                     output_col_dim, &output_shape));
  }
  c->set_output(0, output_shape);
  return Status::OK();
}


Status ConcatShapeHelper(InferenceContext* c, int start_value_index,
                         int end_value_index, int dim_index) {
  ShapeHandle unused;
  TF_RETURN_IF_ERROR(c->WithRank(c->input(dim_index), 0, &unused));
  const Tensor* concat_dim_t = c->input_tensor(dim_index);
  if (concat_dim_t == nullptr) {
    // Return an unknown shape with same rank as inputs, or an unknown rank
    // if no input's rank is known.

    // Find rank.
    int32_t rank = InferenceContext::kUnknownRank;
    for (int i = start_value_index; i < end_value_index; ++i) {
      if (rank == InferenceContext::kUnknownRank) rank = c->Rank(c->input(i));
      if (rank != InferenceContext::kUnknownRank) {
        break;
      }
    }
    if (rank == InferenceContext::kUnknownRank) {
      c->set_output(0, c->UnknownShape());
      return Status::OK();
    } else if (rank == 0) {
      return errors::InvalidArgument(
          "Can't concatenate scalars (use tf.stack instead)");
    } else {
      for (int i = start_value_index; i < end_value_index; ++i) {
        // Check that all the inputs are of the correct rank.
        TF_RETURN_IF_ERROR(c->WithRank(c->input(i), rank, &unused));
      }
    }
    // Build result of <rank> different unknown dims.
    std::vector<DimensionHandle> dims;
    dims.reserve(rank);
    for (int i = 0; i < rank; ++i) dims.push_back(c->UnknownDim());
    c->set_output(0, c->MakeShape(dims));
    return Status::OK();
  }

  // Merge all the non-concat dims, and sum the concat dim to make an output
  // shape.
  int64_t concat_dim;
  if (concat_dim_t->dtype() == DT_INT32) {
    concat_dim = static_cast<int64_t>(concat_dim_t->flat<int32>()(0));
  } else {
    concat_dim = concat_dim_t->flat<int64_t>()(0);
  }

  // Minimum required number of dimensions.
  const int64 min_rank = concat_dim < 0 ? -concat_dim : concat_dim + 1;

  ShapeHandle output_before;
  ShapeHandle output_after;

  ShapeHandle input = c->input(end_value_index - 1);
  TF_RETURN_IF_ERROR(c->WithRankAtLeast(input, min_rank, &input));
  TF_RETURN_IF_ERROR(c->Subshape(input, 0, concat_dim, &output_before));
  DimensionHandle output_middle = c->Dim(input, concat_dim);
  if (concat_dim == -1) {
    output_after = c->Scalar();  // no dimensions.
  } else {
    TF_RETURN_IF_ERROR(c->Subshape(input, concat_dim + 1, &output_after));
  }

  for (int i = end_value_index - 2; i >= start_value_index; --i) {
    ShapeHandle before;
    ShapeHandle after;
    input = c->input(i);
    TF_RETURN_IF_ERROR(c->WithRankAtLeast(input, min_rank, &input));
    TF_RETURN_IF_ERROR(c->Subshape(input, 0, concat_dim, &before));
    DimensionHandle middle = c->Dim(input, concat_dim);
    if (concat_dim == -1) {
      after = c->Scalar();
    } else {
      TF_RETURN_IF_ERROR(c->Subshape(input, concat_dim + 1, &after));
    }

    TF_RETURN_IF_ERROR(c->Merge(before, output_before, &output_before));
    TF_RETURN_IF_ERROR(c->Add(output_middle, middle, &output_middle));
    TF_RETURN_IF_ERROR(c->Merge(after, output_after, &output_after));
  }

  ShapeHandle s;
  TF_RETURN_IF_ERROR(
      c->Concatenate(output_before, c->Vector(output_middle), &s));
  TF_RETURN_IF_ERROR(c->Concatenate(s, output_after, &s));
  c->set_output(0, s);
  return Status::OK();
}


Status Conv2DBackpropInputShape(shape_inference::InferenceContext* c) {
  string data_format_str;
  if (!c->GetAttr("data_format", &data_format_str).ok()) {
    data_format_str = "NHWC";
  }
  TensorFormat data_format;
  if (!FormatFromString(data_format_str, &data_format)) {
    return errors::InvalidArgument("Invalid data format string: ",
                                   data_format_str);
  }

  // For the rest of this function, output_grad_* describes out_backprop and
  // input_grad_* describes in_backprop.
  ShapeHandle output_grad_shape = c->input(2);
  TF_RETURN_IF_ERROR(c->WithRank(output_grad_shape, 4, &output_grad_shape));
  ShapeHandle filter_shape = c->input(1);
  TF_RETURN_IF_ERROR(c->WithRank(filter_shape, 4, &filter_shape));

  DimensionHandle batch_size_dim;
  DimensionHandle output_grad_depth_dim;
  gtl::InlinedVector<DimensionHandle, 2> output_grad_spatial_dims(2);
  TF_RETURN_IF_ERROR(DimensionsFromShape(
      output_grad_shape, data_format, &batch_size_dim,
      absl::MakeSpan(output_grad_spatial_dims), &output_grad_depth_dim, c));
  DimensionHandle unused;
  TF_RETURN_IF_ERROR(
      c->Merge(output_grad_depth_dim, c->Dim(filter_shape, 3), &unused));

  ShapeHandle specified_input_grad_shape;
  TF_RETURN_IF_ERROR(
      c->MakeShapeFromShapeTensor(0, &specified_input_grad_shape));
  if (c->Rank(specified_input_grad_shape) == InferenceContext::kUnknownRank) {
    TF_RETURN_IF_ERROR(c->WithRank(specified_input_grad_shape, 4,
                                   &specified_input_grad_shape));
  }

  // input_grad_depth_dim doesn't equal c->Dim(filter_shape,2) when the number
  // of groups is larger than 1. If input_sizes is a 4D shape, we collect
  // input_grad_depth_dim from input_sizes; otherwise we compute it as
  // c->Dim(filter_shape,2).
  DimensionHandle input_grad_depth_dim;
  gtl::InlinedVector<DimensionHandle, 2> specified_input_grad_spatial_dims(2);
  int specified_input_grad_rank = c->Rank(specified_input_grad_shape);
  if (specified_input_grad_rank == 4) {
    DimensionHandle specified_batch_size_dim;
    TF_RETURN_IF_ERROR(DimensionsFromShape(
        specified_input_grad_shape, data_format, &specified_batch_size_dim,
        absl::MakeSpan(specified_input_grad_spatial_dims),
        &input_grad_depth_dim, c));
    TF_RETURN_IF_ERROR(
        c->Merge(specified_batch_size_dim, batch_size_dim, &unused));
  } else if (specified_input_grad_rank == 2) {
    specified_input_grad_spatial_dims[0] =
        c->Dim(specified_input_grad_shape, 0);
    specified_input_grad_spatial_dims[1] =
        c->Dim(specified_input_grad_shape, 1);
    input_grad_depth_dim = c->Dim(filter_shape, 2);
  } else {
    return errors::InvalidArgument(
        "Conv2DBackpropInput requires input_sizes to contain 4 values or 2 "
        "values, but got: ",
        specified_input_grad_rank);
  }

  ShapeHandle input_grad_shape;
  TF_RETURN_IF_ERROR(ShapeFromDimensions(
      batch_size_dim, specified_input_grad_spatial_dims, input_grad_depth_dim,
      data_format, /*vect_size=*/absl::nullopt, c, &input_grad_shape));
  c->set_output(0, input_grad_shape);
  return Status::OK();
}


Status QuantizeV2Shape(InferenceContext* c) {
  int axis = -1;
  Status s = c->GetAttr("axis", &axis);
  if (!s.ok() && s.code() != error::NOT_FOUND) {
    return s;
  }
  if (axis < -1) {
    return errors::InvalidArgument("axis should be at least -1, got ", axis);
  }
  const int minmax_rank = (axis == -1) ? 0 : 1;
  TF_RETURN_IF_ERROR(shape_inference::UnchangedShape(c));
  ShapeHandle minmax;
  TF_RETURN_IF_ERROR(c->WithRank(c->input(1), minmax_rank, &minmax));
  TF_RETURN_IF_ERROR(c->WithRank(c->input(2), minmax_rank, &minmax));
  if (axis != -1) {
    ShapeHandle input;
    TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), axis + 1, &input));
    DimensionHandle depth;
    TF_RETURN_IF_ERROR(
        c->Merge(c->Dim(minmax, 0), c->Dim(input, axis), &depth));
  }
  c->set_output(1, minmax);
  c->set_output(2, minmax);
  return Status::OK();
}


Status DepthwiseConv2DNativeShapeImpl(shape_inference::InferenceContext* c,
                                      bool supports_explicit_padding) {
  ShapeHandle input_shape;
  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 4, &input_shape));
  ShapeHandle filter_shape;
  TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 4, &filter_shape));

  std::vector<int32> strides;
  TF_RETURN_IF_ERROR(c->GetAttr("strides", &strides));

  if (strides.size() != 4) {
    return errors::InvalidArgument(
        "DepthwiseConv2D requires the stride attribute to contain 4 values, "
        "but got: ",
        strides.size());
  }

  std::vector<int32> dilations;
  if (!c->GetAttr("dilations", &dilations).ok()) {
    dilations.resize(4, 1);
  }

  if (dilations.size() != 4) {
    return errors::InvalidArgument(
        "DepthwiseConv2D requires the dilations attribute to contain 4 values, "
        "but got: ",
        dilations.size());
  }

  string data_format_str;
  Status s = c->GetAttr("data_format", &data_format_str);
  TensorFormat data_format;
  if (!s.ok() || !FormatFromString(data_format_str, &data_format)) {
    data_format = FORMAT_NHWC;
  }
  int32_t stride_rows;
  int32_t stride_cols;
  int32_t dilation_rows;
  int32_t dilation_cols;
  if (data_format == FORMAT_NCHW) {
    // Canonicalize input shape to NHWC so the shape inference code below can
    // process it.
    input_shape =
        c->MakeShape({{c->Dim(input_shape, 0), c->Dim(input_shape, 2),
                       c->Dim(input_shape, 3), c->Dim(input_shape, 1)}});
    stride_rows = strides[2];
    stride_cols = strides[3];
    dilation_rows = dilations[2];
    dilation_cols = dilations[3];
  } else {
    stride_rows = strides[1];
    stride_cols = strides[2];
    dilation_rows = dilations[1];
    dilation_cols = dilations[2];
  }

  DimensionHandle batch_size_dim = c->Dim(input_shape, 0);
  DimensionHandle in_rows_dim = c->Dim(input_shape, 1);
  DimensionHandle in_cols_dim = c->Dim(input_shape, 2);

  DimensionHandle filter_rows_dim = c->Dim(filter_shape, 0);
  DimensionHandle filter_cols_dim = c->Dim(filter_shape, 1);
  DimensionHandle input_depth = c->Dim(filter_shape, 2);
  DimensionHandle depth_multiplier = c->Dim(filter_shape, 3);

  // Check that the input depths are compatible.
  TF_RETURN_IF_ERROR(
      c->Merge(c->Dim(input_shape, 3), input_depth, &input_depth));

  DimensionHandle output_depth;
  TF_RETURN_IF_ERROR(c->Multiply(input_depth, depth_multiplier, &output_depth));

  Padding padding;
  TF_RETURN_IF_ERROR(c->GetAttr("padding", &padding));

  std::vector<int64_t> explicit_paddings;
  if (supports_explicit_padding) {
    Status status = c->GetAttr("explicit_paddings", &explicit_paddings);
    // Use the default value, which is an empty list, if the attribute is not
    // found. Otherwise return the error to the caller.
    if (!status.ok() && !errors::IsNotFound(status)) {
      return status;
    }
    TF_RETURN_IF_ERROR(CheckValidPadding(padding, explicit_paddings,
                                         /*num_dims=*/4, data_format));
  } else {
    DCHECK(padding != Padding::EXPLICIT);
  }

  // TODO(mrry,shlens): Raise an error if the stride would cause
  // information in the input to be ignored. This will require a change
  // in the kernel implementation.
  DimensionHandle output_rows, output_cols;
  int64_t pad_rows_before = -1, pad_rows_after = -1;
  int64_t pad_cols_before = -1, pad_cols_after = -1;
  if (padding == Padding::EXPLICIT) {
    GetExplicitPaddingForDim(explicit_paddings, data_format, 'H',
                             &pad_rows_before, &pad_rows_after);
    GetExplicitPaddingForDim(explicit_paddings, data_format, 'W',
                             &pad_cols_before, &pad_cols_after);
  }
  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDimsV2(
      c, in_rows_dim, filter_rows_dim, dilation_rows, stride_rows, padding,
      pad_rows_before, pad_rows_after, &output_rows));
  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDimsV2(
      c, in_cols_dim, filter_cols_dim, dilation_cols, stride_cols, padding,
      pad_cols_before, pad_cols_after, &output_cols));

  ShapeHandle output_shape;
  if (data_format == FORMAT_NCHW) {
    output_shape =
        c->MakeShape({batch_size_dim, output_depth, output_rows, output_cols});
  } else {
    output_shape =
        c->MakeShape({batch_size_dim, output_rows, output_cols, output_depth});
  }
  c->set_output(0, output_shape);
  return Status::OK();
}


Status FusedBatchNormGradShape(shape_inference::InferenceContext* c) {
  string data_format_str;
  TF_RETURN_IF_ERROR(c->GetAttr("data_format", &data_format_str));
  TensorFormat data_format;
  if (!FormatFromString(data_format_str, &data_format)) {
    return errors::InvalidArgument("Invalid data format string: ",
                                   data_format_str);
  }
  const int rank =
      (data_format_str == "NDHWC" || data_format_str == "NCDHW") ? 5 : 4;
  ShapeHandle y_backprop;
  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), rank, &y_backprop));
  ShapeHandle x;
  TF_RETURN_IF_ERROR(c->WithRank(c->input(1), rank, &x));

  bool is_training;
  TF_RETURN_IF_ERROR(c->GetAttr("is_training", &is_training));

  int channel_dim_index = GetTensorFeatureDimIndex(rank, data_format);
  DimensionHandle channel_dim = c->Dim(y_backprop, channel_dim_index);
  TF_RETURN_IF_ERROR(
      c->Merge(channel_dim, c->Dim(x, channel_dim_index), &channel_dim));

  // covers scale, mean (reserve_space_1), variance (reserve_space_2)
  for (int i = 2; i < 5; ++i) {
    ShapeHandle vec;
    TF_RETURN_IF_ERROR(c->WithRank(c->input(i), 1, &vec));
    TF_RETURN_IF_ERROR(c->Merge(channel_dim, c->Dim(vec, 0), &channel_dim));
  }

  ShapeHandle x_backprop;
  TF_RETURN_IF_ERROR(
      c->ReplaceDim(y_backprop, channel_dim_index, channel_dim, &x_backprop));
  c->set_output(0, x_backprop);
  c->set_output(1, c->Vector(channel_dim));
  c->set_output(2, c->Vector(channel_dim));
  c->set_output(3, c->Vector(0));
  c->set_output(4, c->Vector(0));
  return Status::OK();
}


Status ReduceScatterShape(shape_inference::InferenceContext* c) {
  shape_inference::ShapeHandle in = c->input(0);
  if (!c->RankKnown(in)) {
    // Input shape unknown, so set unknown output shape.
    c->set_output(0, in);
    return Status::OK();
  }

  shape_inference::ShapeHandle group_assignment_shape = c->input(1);
  if (c->Rank(group_assignment_shape) != 2)
    return errors::InvalidArgument(
        "ReduceScatter group_assignment should be rank 2");

  const Tensor* scatter_dimension = c->input_tensor(2);
  if (!scatter_dimension) {
    c->set_output(0, c->UnknownShape());
    return Status::OK();
  }
  int64_t scatter_dim;
  TF_RETURN_IF_ERROR(c->GetScalarFromTensor(scatter_dimension, &scatter_dim));

  std::vector<shape_inference::DimensionHandle> out_dims;
  out_dims.reserve(c->Rank(in));
  for (int i = 0; i < c->Rank(in); ++i) {
    // If the dimension is the scatter_dimension, then divide the dimension
    // by the partition size in the group_assignment.
    if (i == scatter_dim) {
      shape_inference::DimensionHandle dim = c->Dim(in, i);
      shape_inference::DimensionHandle out_dim;
      TF_RETURN_IF_ERROR(c->Divide(dim, c->Dim(group_assignment_shape, 1),
                                   /*evenly_divisible=*/true, &out_dim));
      out_dims.push_back(out_dim);
    } else {
      out_dims.emplace_back(c->Dim(in, i));
    }
  }
  c->set_output(0, c->MakeShape(out_dims));
  return Status::OK();
}