Fix mismatched argument order in TFL->TOSA gather op.

TFL_GatherOp takes in params then indices while the Tosa_GatherOp accepts indices then params. This updates the compiler conversion to swap arguments and adds a test.

PiperOrigin-RevId: 357822827
Change-Id: I538c4ef83367ca7532eef49592019e29accbbc72
This commit is contained in:
A. Unique TensorFlower 2021-02-16 15:52:30 -08:00 committed by TensorFlower Gardener
parent 18b63d903e
commit 80a8cf901f
2 changed files with 9 additions and 2 deletions

View File

@ -81,6 +81,13 @@ func @test_add(%arg0: tensor<13x21x1xf32>, %arg1: tensor<13x21x3xf32>) -> tensor
return %0 : tensor<13x21x3xf32>
}
// CHECK-LABEL: test_gather
// CHECK: tosa.gather
func @test_gather(%arg0: tensor<100x25xf32>, %arg1: tensor<1x20xi32>) -> tensor<20x25x3xf32> {
%0 = "tfl.gather"(%arg0, %arg1) {axis = 0 : i32} : (tensor<100x25xf32>, tensor<1x20xi32>) -> tensor<20x25x3xf32>
return %0 : tensor<20x25x3xf32>
}
// -----
// CHECK-LABEL: test_sub

View File

@ -49,6 +49,6 @@ def : Pat<(TFL_PowOp $l, $r), (Tosa_PowOp $l, $r)>;
def : Pat<(TFL_GatherOp $params,
$indices,
$axis),
(Tosa_GatherOp $params,
$indices,
(Tosa_GatherOp $indices,
$params,
$axis)>;