Optimize array stacking in the output by using NumPy rather than JAX.

PiperOrigin-RevId: 701895415

torax/output.py

@@ -13,14 +13,16 @@
 # limitations under the License.
 
 """Module containing functions for saving and loading simulation output."""
+
 from __future__ import annotations
 
 import dataclasses
 
 from absl import logging
 import chex
 import jax
-from jax import numpy as jnp
+import numpy as np
+
 from torax import state
 from torax.config import runtime_params
 from torax.geometry import geometry
@@ -200,7 +202,11 @@ def __init__(
     post_processed_output = [
         state.post_processed_outputs for state in sim_outputs.sim_history
     ]
-    stack = lambda *ys: jnp.stack(ys)
+
+    def stack(*x):
+      out = np.stack([np.asarray(i) for i in x])
+      return out
+
     self.core_profiles: state.CoreProfiles = jax.tree_util.tree_map(
         stack, *core_profiles
     )
@@ -213,7 +219,7 @@ def __init__(
     self.post_processed_outputs: state.PostProcessedOutputs = (
         jax.tree_util.tree_map(stack, *post_processed_output)
     )
-    self.times = jnp.array([state.t for state in sim_outputs.sim_history])
+    self.times = np.array([state.t for state in sim_outputs.sim_history])
     chex.assert_rank(self.times, 1)
     self.sim_error = sim_outputs.sim_error
     self.source_models = source_models

torax/simulation_app.py

@@ -58,6 +58,7 @@ def run(_):
 from typing import Callable, Final
 
 from absl import logging
+import chex
 import jax
 from torax import output
 from torax import sim as sim_lib
@@ -141,7 +142,7 @@ def _log_single_state(
 def log_simulation_output_to_stdout(
     core_profile_history: state.CoreProfiles,
     geo: geometry.Geometry,
-    t: jax.Array,
+    t: chex.Array,
 ) -> None:
   del geo
   _log_single_state(core_profile_history.index(0), t[0])