Skip to content

Commit

Permalink
topk
Browse files Browse the repository at this point in the history
  • Loading branch information
justinchuby committed Jan 4, 2025
1 parent 135d5fd commit db51048
Showing 2 changed files with 2 additions and 3 deletions.
3 changes: 1 addition & 2 deletions onnxscript/function_libs/torch_lib/ops/core.py
Original file line number Diff line number Diff line change
@@ -8279,8 +8279,7 @@ def aten_topk(
"""topk(Tensor self, int k, int dim=-1, bool largest=True, bool sorted=True) -> (Tensor values, Tensor indices)"""

# We do not handle scalar inputs for topk
k = op.Constant(value_ints=[k])
values, indices = op.TopK(self, k, axis=dim, largest=largest, sorted=sorted)
values, indices = op.TopK(self, [k], axis=dim, largest=largest, sorted=sorted)
return values, indices


2 changes: 1 addition & 1 deletion tests/function_libs/torch_lib/ops_test_data.py
Original file line number Diff line number Diff line change
@@ -1479,7 +1479,7 @@ def _where_input_wrangler(
)
.skip(
matcher=lambda sample: len(sample.input.shape) == 0 or sample.input.numel() == 0,
reason="scalar inputs or empty inputs are not handled"
reason="scalar inputs or empty inputs are not handled",
),
TorchLibOpInfo("tril", core_ops.aten_tril).xfail(
dtypes=(torch.int32,),

0 comments on commit db51048

Please sign in to comment.