Parametrized size passed as constant - how to do it?

[kubaraczkowski]

Hey, thanks or onnxscript! It's really great, I didn't expect this to be so easy and fun!

I'm trying to write a simple percentile function, like e.g. np.percentile(X, 95).
One of the steps (see below) is to have a sorted 1D tensor, for which I use the op.TopK function, since it seems to be the only available 'sorting' in ONNX. Please direct me if I'm wrong about it.

Unfortunately, TopK used as sorting requires the K to be passed as a length of the input tensor X.

I've followed the https://onnxscript.ai/tutorial/index.html#semantics-script-constants, since I know the size of X at 'compile time', (but not necessarily at 'write time of percentile()'. This is my implementation, which gives the following error.

from onnx import helper
X= np.array([5,7,2,8,1,2,3,4,8,9,4,7]).reshape(3,4)
M = helper.make_tensor("M", TensorProto.INT64, (1,), [len(X)])

@script()
def percentile(X: FLOAT[...], p: FLOAT[1])->FLOAT[1]:
    Mc = op.Constant(value=M)
    X_R = op.Reshape(X, [-1])
    X_SORTED, _ = op.TopK(X_R, Mc, axis=0, largest=0)
    IX = p/100.0 * op.Cast(Mc, to=TensorProto.DOUBLE)
    IX_INT = op.Cast(IX, to=TensorProto.INT64)
    return X_SORTED[IX_INT]
print(percentile(X, 90.0))
onnx.checker.check_model(percentile.to_model_proto())

Non-zero status code returned while running TopK node. Name:'' Status Message: k tensor should be a 1D tensor of size 1

I tried using op.Size to get the length of X at runtime, tried different reshapes etc, but still it seems the tensor M (and constant Mc) I construct is not a "1D with 1 element).

Any hints?
Thanks in advance!

[justinchuby]

Consider the following

@script()
def percentile(X: FLOAT[...], p: FLOAT[1])->FLOAT[1]:
    Mc = op.Reshape(op.Size(X), [-1])
    X_R = op.Reshape(X, [-1])
    X_SORTED, _ = op.TopK(X_R, Mc, axis=0, largest=0)
    IX = p/100.0 * op.Cast(Mc, to=TensorProto.DOUBLE)
    IX_INT = op.Cast(IX, to=TensorProto.INT64)
    return X_SORTED[IX_INT]

If you want to sort all of X?

[kubaraczkowski]

That I tried indeed! However, it brings the cryptic:

onnx.onnx_cpp2py_export.checker.ValidationError: Field 'shape' of 'type' is required but missing.

when running onnx.checker.check_model(percentile.to_model_proto())

Additionally - I wonder what the performance implication of using a runtime 'Size()' would be. Since I know the shape beforehand, I don't actually need to compute it. On the other hand, perhaps it gets optimized-out when size of X is fixed?

[justinchuby]

I see.

In X: FLOAT[...], the rank needs to be known for a graph input. I would specify the shape there.

Backends like ONNX Runtime will optimize the model by doing constant folding. If it is known, you may replace the call with a constant too.

[kubaraczkowski]

Right, that does work indeed, even by just giving an 'M' size like this:

@script()
def percentile(X: FLOAT['M'], p: FLOAT[1])->FLOAT[1]:
    Mc = op.Reshape(op.Size(X), [-1])
    print(Mc)
    X_R = op.Reshape(X, [-1])
    X_SORTED, _ = op.TopK(X_R, Mc, axis=0, largest=0)
    IX = p/100.0 * op.Cast(Mc, to=TensorProto.DOUBLE)
    IX_INT = op.Cast(IX, to=TensorProto.INT64)
    return X_SORTED[IX_INT]
onnx.checker.check_model(percentile.to_model_proto())
X= np.arange(1000).reshape(10,10,10)
percentile(X,95.0)

As does this, providing the X.size as a constant:

X= np.arange(1000).reshape(10,10,10)

M = helper.make_tensor("M", TensorProto.INT64, (1,), [X.size])

@script()
def percentile(X: FLOAT['M'], p: FLOAT[1])->FLOAT[1]:
    Mco = op.Constant(value=M)
    Mc = op.Reshape(Mco, [-1])
    X_R = op.Reshape(X, [-1])
    X_SORTED, _ = op.TopK(X_R, Mc, axis=0, largest=0)
    IX = p/100.0 * op.Cast(Mc, to=TensorProto.DOUBLE)
    IX_INT = op.Cast(IX, to=TensorProto.INT64)
    return X_SORTED[IX_INT]
onnx.checker.check_model(percentile.to_model_proto())
percentile(X,95.0)

Thanks a lot for your help @justinchuby , onnxscript rules!