gemini-code-assist[bot] commented on code in PR #19654:
URL: https://github.com/apache/tvm/pull/19654#discussion_r3333024356
##########
python/tvm/relax/frontend/tflite/tflite_frontend.py:
##########
@@ -755,6 +756,88 @@ def convert_hashtable_find(self, op):
"HASHTABLE_FIND requires TensorType.STRING support in Relax TFLite
frontend"
)
+ def convert_hashtable_lookup(self, op):
+ """Convert TFLite HASHTABLE_LOOKUP for non-string value tensors."""
+ from tflite.TensorType import TensorType
+
+ input_tensors = self.get_input_tensors(op)
+ output_tensors = self.get_output_tensors(op)
+ if len(input_tensors) != 3 or len(output_tensors) != 2:
+ raise tvm.error.OpNotImplemented(
+ "HASHTABLE_LOOKUP expects lookup, key, and value inputs with
two outputs"
+ )
+
+ lookup_tensor, key_tensor, value_tensor = input_tensors
+ output_tensor, hits_tensor = output_tensors
+
+ if (
+ lookup_tensor.tensor.Type() != TensorType.INT32
+ or key_tensor.tensor.Type() != TensorType.INT32
+ ):
+ raise tvm.error.OpNotImplemented(
+ "HASHTABLE_LOOKUP requires int32 lookup and key tensors"
+ )
+ if self._is_tflite_string_type(value_tensor.tensor.Type()):
+ raise tvm.error.OpNotImplemented(
+ "HASHTABLE_LOOKUP with TensorType.STRING values is not
supported"
+ )
+ if value_tensor.tensor.Type() != output_tensor.tensor.Type():
+ raise tvm.error.OpNotImplemented(
+ "HASHTABLE_LOOKUP output dtype must match the value tensor
dtype"
+ )
+ if hits_tensor.tensor.Type() != TensorType.UINT8:
+ raise tvm.error.OpNotImplemented("HASHTABLE_LOOKUP hits output
must be uint8")
+
+ lookup_shape = to_int_list(self.get_tensor_shape(lookup_tensor))
+ key_shape = to_int_list(self.get_tensor_shape(key_tensor))
+ value_shape = to_int_list(self.get_tensor_shape(value_tensor))
+ output_shape = to_int_list(self.get_tensor_shape(output_tensor))
+ hits_shape = to_int_list(self.get_tensor_shape(hits_tensor))
+
+ if len(lookup_shape) != 1 or len(key_shape) != 1 or len(value_shape) <
1:
+ raise tvm.error.OpNotImplemented(
+ "HASHTABLE_LOOKUP requires rank-1 lookup/key and rank>=1 value
tensors"
+ )
+ if key_shape[0] != value_shape[0]:
+ raise tvm.error.OpNotImplemented(
+ "HASHTABLE_LOOKUP requires key and value tensors to agree on
row count"
+ )
+ if key_shape[0] == 0:
+ raise tvm.error.OpNotImplemented(
+ "HASHTABLE_LOOKUP requires a non-empty key/value table"
+ )
+ if output_shape != [lookup_shape[0]] + value_shape[1:]:
+ raise tvm.error.OpNotImplemented(
+ "HASHTABLE_LOOKUP output shape must match lookup count and
value tail shape"
+ )
+ if hits_shape != [lookup_shape[0]]:
+ raise tvm.error.OpNotImplemented(
+ "HASHTABLE_LOOKUP hits output shape must match lookup count"
+ )
+
+ lookup = self.get_tensor_expr(lookup_tensor)
+ key = self.get_tensor_expr(key_tensor)
+ value = self.get_tensor_expr(value_tensor)
+
+ positions = relax.op.bucketize(lookup, key, out_int32=True,
right=False)
+ candidate_keys = relax.op.take(key, positions, axis=0, mode="clip")
+ in_range = relax.op.less(positions, relax.const(key_shape[0], "int32"))
+ found = relax.op.logical_and(in_range, relax.op.equal(candidate_keys,
lookup))
Review Comment:
The current implementation uses `relax.op.bucketize` to find the lookup
positions. However, `bucketize` assumes that the boundaries (the `key` tensor)
are sorted. In TFLite, `HASHTABLE_LOOKUP` keys are not guaranteed to be sorted.
If the keys are unsorted, `bucketize` will perform an incorrect binary search
and return wrong indices.
To support unsorted keys correctly, we can use a broadcasted comparison to
find the matching indices using `relax.op.equal`, `relax.op.argmax`, and
`relax.op.max`.
```suggestion
lookup_expanded = relax.op.expand_dims(lookup, axis=1)
key_expanded = relax.op.expand_dims(key, axis=0)
match_matrix = relax.op.equal(lookup_expanded, key_expanded)
match_matrix_int = relax.op.astype(match_matrix, "int32")
positions = relax.op.astype(relax.op.argmax(match_matrix_int,
axis=1), "int32")
found = relax.op.astype(relax.op.max(match_matrix_int, axis=1),
"bool")
```
##########
tests/python/relax/test_frontend_tflite.py:
##########
@@ -5908,6 +5950,53 @@ def test_hashtable_size_uninitialized_unsupported():
_load_model_from_buffer(_build_tflite_hashtable_size_uninitialized_model())
+def test_hashtable_lookup_1d_value():
+ mod =
_load_model_from_buffer(_build_tflite_hashtable_lookup_model(value_shape=[3]))
+
+ output, hits = _run_module(
+ mod,
+ np.array([1234, -292, -11, 0], dtype=np.int32),
+ np.array([-11, 0, 1234], dtype=np.int32),
+ np.array([0.0, 0.1, 0.4], dtype=np.float32),
+ )
+
+ np.testing.assert_allclose(output, np.array([0.4, 0.0, 0.0, 0.1],
dtype=np.float32))
+ np.testing.assert_array_equal(hits, np.array([1, 0, 1, 1], dtype=np.uint8))
Review Comment:
Since TFLite hashtable keys are not guaranteed to be sorted, we should
update the test case to use unsorted keys to ensure correctness and prevent
regressions.
```suggestion
output, hits = _run_module(
mod,
np.array([1234, -292, -11, 0], dtype=np.int32),
np.array([0, 1234, -11], dtype=np.int32),
np.array([0.1, 0.4, 0.0], dtype=np.float32),
)
np.testing.assert_allclose(output, np.array([0.4, 0.0, 0.0, 0.1],
dtype=np.float32))
np.testing.assert_array_equal(hits, np.array([1, 0, 1, 1],
dtype=np.uint8))
```
--
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]
