|
1 | 1 | import functools |
2 | 2 | import re |
3 | 3 |
|
| 4 | +from accelforge.util import NUMPY_FLOAT_TYPE |
4 | 5 | from accelforge.util._frozenset import fzs |
5 | 6 | from accelforge.frontend.workload import Rank |
6 | 7 | from accelforge.util._base_analysis_types import ActionKey, VerboseActionKey |
@@ -208,18 +209,22 @@ def is_fused_loop_col(c: str) -> bool: |
208 | 209 | def is_n_iterations_col(c: str) -> bool: |
209 | 210 | return c.startswith("fused_loop<SEP>n_iterations") |
210 | 211 |
|
211 | | - |
212 | | -def add_to_col(df, target, source): |
| 212 | +def ensure_float_type(df, target, source): |
213 | 213 | if target in df: |
214 | 214 | target_type = df[target].dtype |
215 | 215 | source_type = df[source].dtype |
216 | 216 | if target_type != source_type: |
217 | | - df[target] = df[target].astype("float64") |
218 | | - df[source] = df[source].astype("float64") |
| 217 | + df[target] = df[target].astype(NUMPY_FLOAT_TYPE) |
| 218 | + df[source] = df[source].astype(NUMPY_FLOAT_TYPE) |
| 219 | + |
| 220 | + |
| 221 | +def add_to_col(df, target, source): |
| 222 | + ensure_float_type(df, target, source) |
219 | 223 | df.loc[:, target] = df[target] + df[source] if target in df else df[source] |
220 | 224 |
|
221 | 225 |
|
222 | 226 | def max_to_col(df, target, source): |
| 227 | + ensure_float_type(df, target, source) |
223 | 228 | df.loc[:, target] = df[[target, source]].max(axis=1) if target in df else df[source] |
224 | 229 |
|
225 | 230 |
|
|
0 commit comments