gemini-code-assist[bot] commented on code in PR #19763:
URL: https://github.com/apache/tvm/pull/19763#discussion_r3409647339
##########
python/tvm/relax/frontend/tflite/tflite_frontend.py:
##########
@@ -7580,6 +7584,69 @@ def convert_fake_quant(self, op):
rounded = relax.op.floor(_op.add(_op.multiply(clamped_shifted,
inv_scale), half))
return relax.op.add(_op.multiply(rounded, scale_expr), nudged_min_expr)
+ def convert_real(self, op):
+ """Convert TFLite REAL op.
+
+ TFLite complex64 tensors are represented as float32[..., 2] in Relax,
+ where index 0 = real part, index 1 = imaginary part along the last axis
+ """
+ input_tensors = self.get_input_tensors(op)
+ assert len(input_tensors) == 1, "input tensors length should be 1"
+ input_tensor = self.get_expr(input_tensors[0].tensor_idx)
+ last_axis = int(input_tensor.struct_info.ndim) - 1
+ # slice last axis at index 0, and squeeze to remove the last axis
+ real = _op.strided_slice(input_tensor, begin=[0], end=[1],
strides=[1], axes=[last_axis])
+ return _op.squeeze(real, axis=[last_axis])
Review Comment:
Using `int(input_tensor.struct_info.ndim) - 1` to find the last axis can
cause a `TypeError` or `ValueError` if the input tensor has dynamic or unknown
rank (where `ndim` is `None` or symbolic). Since TVM Relax operators like
`strided_slice` and `squeeze` natively support negative indexing, we can
simplify this and make it more robust by using `axes=[-1]` and `axis=[-1]`
directly.
```suggestion
# slice last axis at index 0, and squeeze to remove the last axis
real = _op.strided_slice(input_tensor, begin=[0], end=[1],
strides=[1], axes=[-1])
return _op.squeeze(real, axis=[-1])
```
##########
tests/python/relax/test_frontend_tflite.py:
##########
@@ -13020,5 +13020,81 @@ def test_unidirectional_sequence_rnn_time_major():
assert tuple(int(d) for d in out_shape) == (batch, time, num_units)
+def test_real():
+ class Real(tf.Module):
+ @tf.function(input_signature=[tf.TensorSpec(shape=(2, 4),
dtype=tf.complex64)])
+ def func(self, x):
+ return tf.math.real(x)
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 4, 2), dtype="float32")) -> R.Tensor((2, 4),
dtype="float32"):
+ R.func_attr({"num_input": 1})
+ with R.dataflow():
+ # slice real part (index 0 along last axis)
+ lv: R.Tensor((2, 4, 1), dtype="float32") = R.strided_slice(
+ x, axes=[2], begin=[0], end=[1], strides=[1]
+ )
+ gv: R.Tensor((2, 4), dtype="float32") = R.squeeze(lv, axis=[2])
Review Comment:
Update the expected IR to match the simplified implementation using negative
indexing (`axes=[-1]` and `axis=[-1]`).
```suggestion
# slice real part (index 0 along last axis)
lv: R.Tensor((2, 4, 1), dtype="float32") = R.strided_slice(
x, axes=[-1], begin=[0], end=[1], strides=[1]
)
gv: R.Tensor((2, 4), dtype="float32") = R.squeeze(lv,
axis=[-1])
```
##########
python/tvm/relax/frontend/tflite/tflite_frontend.py:
##########
@@ -7580,6 +7584,69 @@ def convert_fake_quant(self, op):
rounded = relax.op.floor(_op.add(_op.multiply(clamped_shifted,
inv_scale), half))
return relax.op.add(_op.multiply(rounded, scale_expr), nudged_min_expr)
+ def convert_real(self, op):
+ """Convert TFLite REAL op.
+
+ TFLite complex64 tensors are represented as float32[..., 2] in Relax,
+ where index 0 = real part, index 1 = imaginary part along the last axis
+ """
+ input_tensors = self.get_input_tensors(op)
+ assert len(input_tensors) == 1, "input tensors length should be 1"
+ input_tensor = self.get_expr(input_tensors[0].tensor_idx)
+ last_axis = int(input_tensor.struct_info.ndim) - 1
+ # slice last axis at index 0, and squeeze to remove the last axis
+ real = _op.strided_slice(input_tensor, begin=[0], end=[1],
strides=[1], axes=[last_axis])
+ return _op.squeeze(real, axis=[last_axis])
+
+ def convert_imag(self, op):
+ """Convert TFLite IMAG op.
+
+ See convert_real for representation of complex64 tensors in Relax.
+ """
+ input_tensors = self.get_input_tensors(op)
+ assert len(input_tensors) == 1, "input tensors length should be 1"
+ input_tensor = self.get_expr(input_tensors[0].tensor_idx)
+ last_axis = int(input_tensor.struct_info.ndim) - 1
+ # slice last axis at index 1, and squeeze to remove the last axis
+ imag = _op.strided_slice(input_tensor, begin=[1], end=[2],
strides=[1], axes=[last_axis])
+ return _op.squeeze(imag, axis=[last_axis])
Review Comment:
Similarly to `convert_real`, we can simplify `convert_imag` and make it
robust against dynamic/unknown ranks by using negative indexing (`axes=[-1]`
and `axis=[-1]`) directly.
```suggestion
# slice last axis at index 1, and squeeze to remove the last axis
imag = _op.strided_slice(input_tensor, begin=[1], end=[2],
strides=[1], axes=[-1])
return _op.squeeze(imag, axis=[-1])
```
##########
tests/python/relax/test_frontend_tflite.py:
##########
@@ -13020,5 +13020,81 @@ def test_unidirectional_sequence_rnn_time_major():
assert tuple(int(d) for d in out_shape) == (batch, time, num_units)
+def test_real():
+ class Real(tf.Module):
+ @tf.function(input_signature=[tf.TensorSpec(shape=(2, 4),
dtype=tf.complex64)])
+ def func(self, x):
+ return tf.math.real(x)
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 4, 2), dtype="float32")) -> R.Tensor((2, 4),
dtype="float32"):
+ R.func_attr({"num_input": 1})
+ with R.dataflow():
+ # slice real part (index 0 along last axis)
+ lv: R.Tensor((2, 4, 1), dtype="float32") = R.strided_slice(
+ x, axes=[2], begin=[0], end=[1], strides=[1]
+ )
+ gv: R.Tensor((2, 4), dtype="float32") = R.squeeze(lv, axis=[2])
+ R.output(gv)
+ return gv
+
+ verify(Real, Expected)
+
+
+def test_imag():
+ class Imag(tf.Module):
+ @tf.function(input_signature=[tf.TensorSpec(shape=(2, 4),
dtype=tf.complex64)])
+ def func(self, x):
+ return tf.math.imag(x)
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 4, 2), dtype="float32")) -> R.Tensor((2, 4),
dtype="float32"):
+ R.func_attr({"num_input": 1})
+ with R.dataflow():
+ # slice imaginary part (index 1 along last axis)
+ lv: R.Tensor((2, 4, 1), dtype="float32") = R.strided_slice(
+ x, axes=[2], begin=[1], end=[2], strides=[1]
+ )
+ gv: R.Tensor((2, 4), dtype="float32") = R.squeeze(lv, axis=[2])
+ R.output(gv)
+ return gv
+
+ verify(Imag, Expected)
+
+
+def test_complex_abs():
+ class ComplexAbs(tf.Module):
+ @tf.function(input_signature=[tf.TensorSpec(shape=(2, 4),
dtype=tf.complex64)])
+ def func(self, x):
+ return tf.math.abs(x)
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 4, 2), dtype="float32")) -> R.Tensor((2, 4),
dtype="float32"):
+ R.func_attr({"num_input": 1})
+ with R.dataflow():
+ lv0: R.Tensor((2, 4, 1), dtype="float32") = R.strided_slice(
+ x, axes=[2], begin=[0], end=[1], strides=[1]
+ )
+ real: R.Tensor((2, 4), dtype="float32") = R.squeeze(lv0,
axis=[2])
+ lv1: R.Tensor((2, 4, 1), dtype="float32") = R.strided_slice(
+ x, axes=[2], begin=[1], end=[2], strides=[1]
+ )
+ imag: R.Tensor((2, 4), dtype="float32") = R.squeeze(lv1,
axis=[2])
Review Comment:
Update the expected IR to match the simplified implementation using negative
indexing (`axes=[-1]` and `axis=[-1]`).
```suggestion
lv0: R.Tensor((2, 4, 1), dtype="float32") = R.strided_slice(
x, axes=[-1], begin=[0], end=[1], strides=[1]
)
real: R.Tensor((2, 4), dtype="float32") = R.squeeze(lv0,
axis=[-1])
lv1: R.Tensor((2, 4, 1), dtype="float32") = R.strided_slice(
x, axes=[-1], begin=[1], end=[2], strides=[1]
)
imag: R.Tensor((2, 4), dtype="float32") = R.squeeze(lv1,
axis=[-1])
```
##########
tests/python/relax/test_frontend_tflite.py:
##########
@@ -13020,5 +13020,81 @@ def test_unidirectional_sequence_rnn_time_major():
assert tuple(int(d) for d in out_shape) == (batch, time, num_units)
+def test_real():
+ class Real(tf.Module):
+ @tf.function(input_signature=[tf.TensorSpec(shape=(2, 4),
dtype=tf.complex64)])
+ def func(self, x):
+ return tf.math.real(x)
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 4, 2), dtype="float32")) -> R.Tensor((2, 4),
dtype="float32"):
+ R.func_attr({"num_input": 1})
+ with R.dataflow():
+ # slice real part (index 0 along last axis)
+ lv: R.Tensor((2, 4, 1), dtype="float32") = R.strided_slice(
+ x, axes=[2], begin=[0], end=[1], strides=[1]
+ )
+ gv: R.Tensor((2, 4), dtype="float32") = R.squeeze(lv, axis=[2])
+ R.output(gv)
+ return gv
+
+ verify(Real, Expected)
+
+
+def test_imag():
+ class Imag(tf.Module):
+ @tf.function(input_signature=[tf.TensorSpec(shape=(2, 4),
dtype=tf.complex64)])
+ def func(self, x):
+ return tf.math.imag(x)
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 4, 2), dtype="float32")) -> R.Tensor((2, 4),
dtype="float32"):
+ R.func_attr({"num_input": 1})
+ with R.dataflow():
+ # slice imaginary part (index 1 along last axis)
+ lv: R.Tensor((2, 4, 1), dtype="float32") = R.strided_slice(
+ x, axes=[2], begin=[1], end=[2], strides=[1]
+ )
+ gv: R.Tensor((2, 4), dtype="float32") = R.squeeze(lv, axis=[2])
Review Comment:
Update the expected IR to match the simplified implementation using negative
indexing (`axes=[-1]` and `axis=[-1]`).
```suggestion
# slice imaginary part (index 1 along last axis)
lv: R.Tensor((2, 4, 1), dtype="float32") = R.strided_slice(
x, axes=[-1], begin=[1], end=[2], strides=[1]
)
gv: R.Tensor((2, 4), dtype="float32") = R.squeeze(lv,
axis=[-1])
```
##########
python/tvm/relax/frontend/tflite/tflite_frontend.py:
##########
@@ -7580,6 +7584,69 @@ def convert_fake_quant(self, op):
rounded = relax.op.floor(_op.add(_op.multiply(clamped_shifted,
inv_scale), half))
return relax.op.add(_op.multiply(rounded, scale_expr), nudged_min_expr)
+ def convert_real(self, op):
+ """Convert TFLite REAL op.
+
+ TFLite complex64 tensors are represented as float32[..., 2] in Relax,
+ where index 0 = real part, index 1 = imaginary part along the last axis
+ """
+ input_tensors = self.get_input_tensors(op)
+ assert len(input_tensors) == 1, "input tensors length should be 1"
+ input_tensor = self.get_expr(input_tensors[0].tensor_idx)
+ last_axis = int(input_tensor.struct_info.ndim) - 1
+ # slice last axis at index 0, and squeeze to remove the last axis
+ real = _op.strided_slice(input_tensor, begin=[0], end=[1],
strides=[1], axes=[last_axis])
+ return _op.squeeze(real, axis=[last_axis])
+
+ def convert_imag(self, op):
+ """Convert TFLite IMAG op.
+
+ See convert_real for representation of complex64 tensors in Relax.
+ """
+ input_tensors = self.get_input_tensors(op)
+ assert len(input_tensors) == 1, "input tensors length should be 1"
+ input_tensor = self.get_expr(input_tensors[0].tensor_idx)
+ last_axis = int(input_tensor.struct_info.ndim) - 1
+ # slice last axis at index 1, and squeeze to remove the last axis
+ imag = _op.strided_slice(input_tensor, begin=[1], end=[2],
strides=[1], axes=[last_axis])
+ return _op.squeeze(imag, axis=[last_axis])
+
+ def convert_complex_abs(self, op):
+ """Convert TFLite COMPLEX_ABS op: sqrt(real^2 + imag^2)
+
+ See convert_real for the float32[..., 2] complex representation
convention.
+ """
+ input_tensors = self.get_input_tensors(op)
+ assert len(input_tensors) == 1, "input tensors length should be 1"
+ input_tensor = self.get_expr(input_tensors[0].tensor_idx)
+ last_axis = int(input_tensor.struct_info.ndim) - 1
+ real = self.bb.emit(
+ _op.strided_slice(input_tensor, begin=[0], end=[1], strides=[1],
axes=[last_axis])
+ )
+ real = self.bb.emit(_op.squeeze(real, axis=[last_axis]))
+ imag = self.bb.emit(
+ _op.strided_slice(input_tensor, begin=[1], end=[2], strides=[1],
axes=[last_axis])
+ )
+ imag = self.bb.emit(_op.squeeze(imag, axis=[last_axis]))
Review Comment:
Similarly to `convert_real`, we can simplify `convert_complex_abs` and make
it robust against dynamic/unknown ranks by using negative indexing (`axes=[-1]`
and `axis=[-1]`) directly.
```suggestion
real = self.bb.emit(
_op.strided_slice(input_tensor, begin=[0], end=[1], strides=[1],
axes=[-1])
)
real = self.bb.emit(_op.squeeze(real, axis=[-1]))
imag = self.bb.emit(
_op.strided_slice(input_tensor, begin=[1], end=[2], strides=[1],
axes=[-1])
)
imag = self.bb.emit(_op.squeeze(imag, axis=[-1]))
```
--
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]
