Skip to content

Commit d367375

Browse files
committed
Simplify LiteRT test utilities using keras ops and standardize_dtype
Replace manual tensor->numpy and dtype conversion chains in: - _build_litert_torch_input_signature: use ops.convert_to_numpy + standardize_dtype - convert_for_tflite: use ops.convert_to_numpy + standardize_dtype + dtype_map Reduces duplication and unifies on existing keras APIs.
1 parent 1e0347a commit d367375

1 file changed

Lines changed: 19 additions & 35 deletions

File tree

keras_hub/src/tests/test_case.py

Lines changed: 19 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -584,22 +584,16 @@ def _build_litert_torch_input_signature(input_data):
584584
The torch export path does not support dynamic shapes, so it needs a
585585
fully specified `keras.InputSpec` tree derived from the sample data.
586586
"""
587-
588-
def _to_numpy(x):
589-
if hasattr(x, "detach"):
590-
return x.detach().cpu().numpy()
591-
if hasattr(x, "numpy") and not isinstance(x, np.ndarray):
592-
return x.numpy()
593-
return x
587+
dtype_map = {
588+
"float64": "float32",
589+
"int64": "int32",
590+
}
594591

595592
def _to_spec(x):
596-
x = _to_numpy(x)
597-
dtype = np.dtype(x.dtype)
598-
if dtype == np.dtype("float64"):
599-
dtype = np.dtype("float32")
600-
elif dtype == np.dtype("int64"):
601-
dtype = np.dtype("int32")
602-
return keras.InputSpec(shape=x.shape, dtype=dtype.name)
593+
x = ops.convert_to_numpy(x)
594+
dtype = keras.backend.standardize_dtype(x.dtype)
595+
dtype = dtype_map.get(dtype, dtype)
596+
return keras.InputSpec(shape=x.shape, dtype=dtype)
603597

604598
return [tree.map_structure(_to_spec, input_data)]
605599

@@ -840,29 +834,19 @@ def run_litert_export_test(
840834
runner = interpreter.get_signature_runner("serving_default")
841835

842836
# Convert input data dtypes to match TFLite expectations
837+
dtype_map = {
838+
"bool": "int32",
839+
"float64": "float32",
840+
"int64": "int32",
841+
}
842+
843843
def convert_for_tflite(x):
844844
"""Convert tensor/array to TFLite-compatible dtypes."""
845-
if hasattr(x, "detach"):
846-
x = x.detach().cpu().numpy()
847-
if hasattr(x, "dtype"):
848-
if isinstance(x, np.ndarray):
849-
if x.dtype == bool:
850-
return x.astype(np.int32)
851-
elif x.dtype == np.float64:
852-
return x.astype(np.float32)
853-
elif x.dtype == np.int64:
854-
return x.astype(np.int32)
855-
else: # TensorFlow tensor
856-
if x.dtype == tf.bool:
857-
return ops.cast(x, "int32").numpy()
858-
elif x.dtype == tf.float64:
859-
return ops.cast(x, "float32").numpy()
860-
elif x.dtype == tf.int64:
861-
return ops.cast(x, "int32").numpy()
862-
else:
863-
return x.numpy() if hasattr(x, "numpy") else x
864-
elif hasattr(x, "numpy"):
865-
return x.numpy()
845+
x = ops.convert_to_numpy(x)
846+
dtype = keras.backend.standardize_dtype(x.dtype)
847+
target = dtype_map.get(dtype)
848+
if target is not None:
849+
x = x.astype(target)
866850
return x
867851

868852
if isinstance(input_data, dict):

0 commit comments

Comments
 (0)