Add specification for SegmentMax-16

[p-wysocki]

Details:

• Specification for tf.math.segment_max (https://www.tensorflow.org/api_docs/python/tf/math/segment_max)

Related PRs:

• Add SegmentMax-16 to Core #28698

Tickets:

• CVS-158914

[rkazants]

Let us have EmbeddingSegmentsMax similar to EmbeddingSegmentsSum.
It should also have default index (defining default value for empty segment)

[rkazants]

please define input shapes and output shape for each input and describe what dimensions are equal

[p-wysocki]

data can have any rank and dimensions, so it's described as ND of any numerical type. segment_ids are specified to be a 1D tensor of non-negative, sorted integer numbers of size equal to the size of the first dimension of the input tensor.

Could you please specify what's missing? I think the shapes are covered, but I may be missing something.

[p-wysocki]

Let us have EmbeddingSegmentsMax similar to EmbeddingSegmentsSum.

EmbeddingSegmentX is for sparse inputs, while SegmentMax I'm implementing (to unlock some models) accepts dense inputs, so I don't think it should be added as EmbeddingSegmentMax. I'm moving the discussion to internal channels, if it results in changes, I'll apply them to the PR.

[mitruska]

I can see that unsorted segment_ids may lead to error or undefined behavior (implementation specific, depends on the hardware).
Should we specify a common behavior for OV op?
Can be clarified at the plugin implementation stage.

[p-wysocki]

I assumed that we'll verify the orderliness of segment_ids during shape inference and raise an error if needed. We can always remove that check from shape inference later and unify the behavior, but I suppose for now the op will be enabled in CPU with reusage of reference implementation.

[p-wysocki]

There is also V2 https://www.tensorflow.org/api_docs/python/tf/raw_ops/SegmentMaxV2 where the default has been changed to numeric_limits::lowest(). Adding attribute for default value seems to be a simple solution to support both cases, but to enable V2 at once we would also need to consider "num_segments" input.

I added an attribute for the default value. num_segments, according to TF docs, is only used for shape inference so it shouldn't be much more work.

[rkazants]

let us have it as an input. it will be more convenient if this value is not constant to compute.

[p-wysocki]

After some discussion with @mitruska I modified the empty_segment_value to still be an attribute, but now is called fill_mode and is of type string. The reason is that it's not an input to the TF operation, but a way to differentiate between SegmentMax and SegmentMaxV2. It allows us to create one operator and enable both of these.

The reason why its type has been changed to string is that the value LOWEST inserts the lower limit of type T into empty segments, and that lower value should be calculated at runtime because we may not know the type during operator conversion and some precision transformations may break the op.

What do you think? Can it stay as an attribute in its current form?

[rkazants]

Let it be a constant input (the fourth one) of type T please. It will be more flexible for future use cases.
Attribute is injected in IR and no way to change it in run-time

[p-wysocki]

Since one of the values is LOWEST, which means the lowest value of type T, wouldn't this require custom behavior for graph type conversions? Setting that value dynamically during runtime could be safer.

[rkazants]

I don't know but I don't exclude that this operation can be used in some decompositions to express some new FW operation. If we make an attribute for it, we create additional constraints. I think we don't need addtional constraints and let it make more flexible.

[rkazants]

still propose to avoid the attribute that can be input. Attribute is a constraint