Apply suggestions from code review

Co-authored-by: Marcel Rath <75042654+ratheron@users.noreply.github.com>

Author: amacati

crazyflow/sim/sharding.py

@@ -73,7 +73,7 @@ def shard(data: SimData, mesh: Mesh) -> SimData:
     return jax.device_put(data, placement(data, mesh))
 
 
-def build_sharded(
+def build_sharded_data(
     create: Callable[[int | Array], SimData], rng_key: int | Array, mesh: Mesh
 ) -> SimData:
     """Build simulation data distributed over a mesh.

crazyflow/sim/sim.py

@@ -35,7 +35,7 @@
 from crazyflow.sim.data import SimControls, SimCore, SimData, SimParams, SimState, SimStateDeriv
 from crazyflow.sim.integration import Integrator, euler, rk4, symplectic_euler
 from crazyflow.sim.pipeline import append_fn
-from crazyflow.sim.sharding import build_sharded, placement
+from crazyflow.sim.sharding import build_sharded_data, placement
 from crazyflow.utils import grid_2d, pytree_replace, world_mask
 
 if TYPE_CHECKING:
@@ -143,7 +143,7 @@ def __init__(
         if mesh is None:
             self.data = self.init_data(*freqs, rng_key)
         else:
-            self.data = build_sharded(partial(self.init_data, *freqs), rng_key, mesh)
+            self.data = build_sharded_data(partial(self.init_data, *freqs), rng_key, mesh)
         self.default_data: SimData = self.build_default_data()
 
         # Build the simulation pipeline and overwrite the default _step implementation with it
@@ -454,7 +454,7 @@ def build_data(self) -> SimData:
         if self.mesh is None:
             self.data = self.init_data(*freqs, rng_key)
         else:
-            self.data = build_sharded(partial(self.init_data, *freqs), rng_key, self.mesh)
+            self.data = build_sharded_data(partial(self.init_data, *freqs), rng_key, self.mesh)
         return self.data
 
     def shard(self, mesh: Mesh) -> SimData: