Hi,

I’m trying to implement a lexicographic max operator that does the following:

given a tensor of N dimensional elements, return the first non-zero value according to the given order.

so for example:

input = [[0, 1, 2],

[1, 4, 2],

[0, 0, 3],

[0, 3, 1]]

output = [[1],

[1],

[3],

[3]]

is there any existing method to do so? if not, could you suggest a way to efficiently implement such a thing (that would also work for “batch” inputs).

Thanks!