jax-hpc-profiler 0.2.1__py3-none-any.whl → 0.2.3__py3-none-any.whl

jax_hpc_profiler/timer.py

@@ -28,6 +28,15 @@ class Timer:
             return None
         return cost_analysis[0]['flops']
 
+    def _normalize_memory_units(self, memory_analysis) -> str:  
+        
+        sizes_str = ['B', 'KB', 'MB', 'GB', 'TB' , 'PB']
+        factors = [1 , 1024 , 1024**2 , 1024**3 , 1024**4 , 1024**5]
+        factor = int(np.log10(memory_analysis) // 3)
+        
+        return f"{memory_analysis / factors[factor]:.2f} {sizes_str[factor]}"
+
+    
     def _read_memory_analysis(self, memory_analysis: Any) -> Tuple:
         if memory_analysis is None:
             return None, None, None, None
@@ -38,7 +47,7 @@ class Timer:
 
     def chrono_jit(self, fun: Callable, *args, ndarray_arg=None) -> np.ndarray:
         start = time.perf_counter()
-        out = jax.jit(fun)(*args)
+        out = fun(*args)
         if ndarray_arg is None:
             out.block_until_ready()
         else:
@@ -59,10 +68,13 @@ class Timer:
         self.compiled_code["LOWERED"] = lowered.as_text()
         self.compiled_code["COMPILED"] = compiled.as_text()
         self.profiling_data["FLOPS"] = cost_analysis
-        self.profiling_data["generated_code"] = memory_analysis[0]
-        self.profiling_data["argument_size"] = memory_analysis[0]
-        self.profiling_data["output_size"] = memory_analysis[0]
-        self.profiling_data["temp_size"] = memory_analysis[0]
+        self.profiling_data[
+            "generated_code"] = memory_analysis[0]
+        self.profiling_data[
+            "argument_size"] = memory_analysis[1]
+        self.profiling_data[
+            "output_size"] = memory_analysis[2]
+        self.profiling_data["temp_size"] = memory_analysis[3]
         return out
 
     def chrono_fun(self, fun: Callable, *args, ndarray_arg=None) -> np.ndarray:
@@ -89,7 +101,7 @@ class Timer:
         global_times = jax.make_array_from_callback(
             shape=global_shape,
             sharding=sharding,
-            data_callback=lambda x: times_array)
+            data_callback=lambda _: jnp.expand_dims(times_array,axis=0))
 
         @partial(shard_map,
                  mesh=mesh,
@@ -101,7 +113,7 @@ class Timer:
 
         times_array = get_mean_times(global_times)
         times_array.block_until_ready()
-        return np.array(times_array.addressable_data(0))
+        return np.array(times_array.addressable_data(0)[0])
 
     def report(self,
                csv_filename: str,
@@ -122,8 +134,6 @@ class Timer:
                 csv_filename), os.path.splitext(
                     os.path.basename(csv_filename))[0]
             report_folder = filename if dirname == "" else f"{dirname}/{filename}"
-            print(
-                f"report_folder: {report_folder} csv_filename: {csv_filename}")
             os.makedirs(report_folder, exist_ok=True)
             md_filename = f"{report_folder}/{x}_{px}_{py}_{backend}_{precision}_{function}.md"
 
@@ -131,13 +141,13 @@ class Timer:
         z = x if z is None else z
 
         times_array = self._get_mean_times()
-
         min_time = np.min(times_array)
         max_time = np.max(times_array)
         mean_time = np.mean(times_array)
         std_time = np.std(times_array)
         last_time = times_array[-1]
 
+
         flops = self.profiling_data["FLOPS"]
         generated_code = self.profiling_data["generated_code"]
         argument_size = self.profiling_data["argument_size"]
@@ -178,12 +188,13 @@ class Timer:
             "Temporary Size": temp_size,
             "FLOPS": self.profiling_data["FLOPS"]
         }
+        iteration_runs = {}
+        for i in range(len(times_array)):
+            iteration_runs[f"Run {i}"] = times_array[i]
 
         with open(md_filename, 'w') as f:
             f.write(f"# Reporting for {function}\n")
             f.write(f"## Parameters\n")
-            keys = list(param_dict.keys())
-            values = list(param_dict.values())
             f.write(
                 tabulate(param_dict.items(),
                          headers=["Parameter", "Value"],
@@ -195,6 +206,12 @@ class Timer:
                          headers=["Parameter", "Value"],
                          tablefmt='github'))
             f.write("\n---\n")
+            f.write(f"## Iteration Runs\n")
+            f.write(
+                tabulate(iteration_runs.items(),
+                         headers=["Iteration", "Time"],
+                         tablefmt='github'))
+            f.write("\n---\n")
             f.write(f"## Compiled Code\n")
             f.write(f"```hlo\n")
             f.write(self.compiled_code["COMPILED"])
@@ -210,3 +227,4 @@ class Timer:
                 f.write(f"```haskel\n")
                 f.write(self.compiled_code["JAXPR"])
                 f.write(f"\n```\n")
+

{jax_hpc_profiler-0.2.1.dist-info → jax_hpc_profiler-0.2.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jax_hpc_profiler
-Version: 0.2.1
+Version: 0.2.3
 Summary: HPC Plotter and profiler for benchmarking data made for JAX
 Author: Wassim Kabalan
 License: GNU GENERAL PUBLIC LICENSE

{jax_hpc_profiler-0.2.1.dist-info → jax_hpc_profiler-0.2.3.dist-info}/RECORD

@@ -2,11 +2,11 @@ jax_hpc_profiler/__init__.py,sha256=yDWt2S0xJZeS6YLBFvyPj1p5txCgFx2kCxJzVZZcdtI,
 jax_hpc_profiler/create_argparse.py,sha256=sY3OKe6lMrXtVnKyx-EtREXLy9L1TK_mdf0WYRQXu5A,6351
 jax_hpc_profiler/main.py,sha256=CKsKVUKsMRatlYfrFLOV1WZ582rZPtofV89sY_2tpQI,2370
 jax_hpc_profiler/plotting.py,sha256=iA-CCBsDK5wM872tmpS6qo4ws9F6MZLutfkdAN_sWlw,8326
-jax_hpc_profiler/timer.py,sha256=4XGKuP2fclGfac2sNz_W8aOamFw7TfiT2Nvp6BarMJk,7621
+jax_hpc_profiler/timer.py,sha256=baE5DRsQBYRBiphkceTi4qI_8FPGKQEh73f2pAeS-oc,8208
 jax_hpc_profiler/utils.py,sha256=GqlGaD-Zf9GmtuQRJAF6S6QD4E2S8QtZ-jGjjC5ZFwU,14133
-jax_hpc_profiler-0.2.1.dist-info/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
-jax_hpc_profiler-0.2.1.dist-info/METADATA,sha256=smuVIDzcbI2aH4pip8Rnh0qsTNjsnVkP8kvCWA1WTWw,49250
-jax_hpc_profiler-0.2.1.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
-jax_hpc_profiler-0.2.1.dist-info/entry_points.txt,sha256=_cFlxSINscX3ZyNiklfjyOOO7vNkddhoYy_v1JQHSO4,51
-jax_hpc_profiler-0.2.1.dist-info/top_level.txt,sha256=DKAhVKDwkerhth-xo7oKFSnnKE0Xm46m94b06vZksA4,17
-jax_hpc_profiler-0.2.1.dist-info/RECORD,,
+jax_hpc_profiler-0.2.3.dist-info/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
+jax_hpc_profiler-0.2.3.dist-info/METADATA,sha256=myC-zD7y_pRb_-tZoSFi0KmglZH8Gk88_-U5RE14Q04,49250
+jax_hpc_profiler-0.2.3.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
+jax_hpc_profiler-0.2.3.dist-info/entry_points.txt,sha256=_cFlxSINscX3ZyNiklfjyOOO7vNkddhoYy_v1JQHSO4,51
+jax_hpc_profiler-0.2.3.dist-info/top_level.txt,sha256=DKAhVKDwkerhth-xo7oKFSnnKE0Xm46m94b06vZksA4,17
+jax_hpc_profiler-0.2.3.dist-info/RECORD,,