-
Notifications
You must be signed in to change notification settings - Fork 30
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
# Links Addresses this issue: [(shortfin) Need debug-print support for f16 arrays and other types not supported array.array types](#292) # Description This PR adds support for directly converting device_arrays to numpy arrays. We also implement some common functions to be used in the debug context, to achieve better visibility into the values of a device_array at runtime. We implement the [ArrayProtocol](https://numpy.org/doc/stable/reference/arrays.interface.html) when numpy is available to make casting between device_arrays and numpy arrays easier. We define a module named `shortfin.array.nputils` to access these features, which is only available in public API if numpy is installed. # Supported Functions - debug_dump_array: Dump contents of array to debug logs - debug_fill_array: Convert to np array, and fill with specific values - _find_mode: Find the mode of an np.array, along a specified dimension - debug_log_tensor_stats: Log the following stats for a tensor: - NaN Count - Shape & dtype - Min - Max - Mode - First 10 Elements - Last 10 Elements --------- Co-authored-by: Xida Ren (Cedar) <[email protected]> Co-authored-by: Stella Laurenzo <[email protected]>
- Loading branch information
3 people
authored
Oct 30, 2024
1 parent
16a1bea
commit 5c6ba0e
Showing
4 changed files
with
441 additions
and
5 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,104 @@ | ||
| import logging | ||
|
|
||
| import numpy as np | ||
|
|
||
| from shortfin import array as sfnp | ||
|
|
||
| logger = logging.getLogger(__name__) | ||
| logger.setLevel(logging.INFO) | ||
|
|
||
|
|
||
| def debug_dump_array(tensor: sfnp.device_array) -> None: | ||
| """Dump the contents of a device array to the debug log. | ||
| Args: | ||
| tensor (sfnp.device_array): The device array to dump. | ||
| """ | ||
| np_array = np.array(tensor) | ||
| logger.debug(np_array) | ||
|
|
||
|
|
||
| def debug_fill_array(tensor: sfnp.device_array, fill_value: int | float) -> np.ndarray: | ||
| """Fill a device array with a given value and return the resulting numpy array. | ||
| Args: | ||
| tensor (sfnp.device_array): The device array to fill. | ||
| fill_value (int | float): The value to fill the array with. | ||
| Returns: | ||
| np.ndarray: The filled numpy array. | ||
| """ | ||
| np_array = np.array(tensor) | ||
| np_array.fill(fill_value) | ||
| return np_array | ||
|
|
||
|
|
||
| def _find_mode( | ||
| arr: np.ndarray, axis=0, keepdims=False | ||
| ) -> tuple[np.ndarray, np.ndarray]: | ||
| """ | ||
| Find the mode of an array along a given axis. | ||
| Args: | ||
| arr: The input array. | ||
| axis: The axis along which to find the mode. | ||
| keepdims: If True, the output shape is the same as arr except along the specified axis. | ||
| Returns: | ||
| tuple: A tuple containing the mode values and the count of the mode values. | ||
| """ | ||
|
|
||
| def _mode(arr): | ||
| if arr.size == 0: | ||
| return np.nan, 0 | ||
|
|
||
| unique, counts = np.unique(arr, return_counts=True) | ||
| max_counts = counts.max() | ||
|
|
||
| mode = unique[counts == max_counts][0] | ||
| return mode, max_counts | ||
|
|
||
| result = np.apply_along_axis(_mode, axis, arr) | ||
| mode_values, mode_count = result[..., 0], result[..., 1] | ||
|
|
||
| if keepdims: | ||
| mode_values = np.expand_dims(mode_values, axis) | ||
| mode_count = np.expand_dims(mode_count, axis) | ||
|
|
||
| return mode_values, mode_count | ||
|
|
||
|
|
||
| def debug_log_tensor_stats(tensor: sfnp.device_array) -> None: | ||
| """Log statistics about a device array to the debug log. | ||
| The following statistics are logged: | ||
| - NaN count | ||
| - Shape, dtype | ||
| - Min, max, mean, mode (excluding NaN values) | ||
| - First 10 elements | ||
| - Last 10 elements | ||
| Args: | ||
| tensor (sfnp.device_array): The device array to log statistics for. | ||
| """ | ||
|
|
||
| np_array = np.array(tensor) | ||
|
|
||
| nan_count = np.isnan(np_array).sum() | ||
|
|
||
| # Remove NaN values | ||
| np_array_no_nan = np_array[~np.isnan(np_array)] | ||
|
|
||
| logger.debug(f"NaN count: {nan_count} / {np_array.size}") | ||
| logger.debug(f"Shape: {np_array.shape}, dtype: {np_array.dtype}") | ||
|
|
||
| if len(np_array_no_nan) > 0: | ||
| mode = _find_mode(np_array_no_nan)[0] | ||
| logger.debug(f"Min (excluding NaN): {np_array_no_nan.min()}") | ||
| logger.debug(f"Max (excluding NaN): {np_array_no_nan.max()}") | ||
| logger.debug(f"Mean (excluding NaN): {np_array_no_nan.mean()}") | ||
| logger.debug(f"Mode (excluding NaN): {mode}") | ||
| logger.debug(f"First 10 elements: {np_array_no_nan.flatten()[:10]}") | ||
| logger.debug(f"Last 10 elements: {np_array_no_nan.flatten()[-10:]}") | ||
| else: | ||
| logger.warning(f"All values are NaN") |
Oops, something went wrong.