gemini-code-assist[bot] commented on code in PR #19516:
URL: https://github.com/apache/tvm/pull/19516#discussion_r3200467992
##########
tests/python/relax/test_frontend_tflite.py:
##########
@@ -1451,6 +1451,64 @@ def main(x: R.Tensor((2, 3), dtype="float32")) ->
R.Tensor((2, 3), dtype="float3
verify(ReverseV2, Expected)
+
+def test_gather():
Review Comment:
This test covers `gather` with `int64` indices. For completeness, it would
be beneficial to also add a test case for `int32` indices, as TFLite's `GATHER`
op supports both `int32` and `int64` for indices.
##########
tests/python/relax/test_frontend_tflite.py:
##########
@@ -1451,6 +1451,64 @@ def main(x: R.Tensor((2, 3), dtype="float32")) ->
R.Tensor((2, 3), dtype="float3
verify(ReverseV2, Expected)
+
+def test_gather():
+ class Gather(tf.Module):
+ @tf.function(
+ input_signature=[
+ tf.TensorSpec(shape=(2, 3, 4), dtype=tf.float32),
+ tf.TensorSpec(shape=(2,), dtype=tf.int64),
+ ]
+ )
+ def func(self, x, indices):
+ return tf.gather(x, indices, axis=1)
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(
+ x: R.Tensor((2, 3, 4), dtype="float32"),
+ indices: R.Tensor((2,), dtype="int64"),
+ ) -> R.Tensor((2, 2, 4), dtype="float32"):
+ R.func_attr({"num_input": 2})
+ with R.dataflow():
+ lv: R.Tensor((2,), dtype="int32") = R.astype(indices,
dtype="int32")
+ gv: R.Tensor((2, 2, 4), dtype="float32") = R.take(x, lv,
axis=1, mode="fast")
+ R.output(gv)
+ return gv
+
+ verify(Gather, Expected)
+
+
+def test_gather_nd():
Review Comment:
This test covers `gather_nd` with `int32` indices, which is great for
verifying the change in this PR. To make the test coverage for this operator
more complete, could you also add a test case for `int64` indices? This would
ensure the `int64` path is not broken by future changes and remains tested.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
