refactoring

Author: priyakasimbeg

algoperf/workloads/criteo1tb/criteo1tb_jax/workload.py

@@ -106,7 +106,7 @@ def init_model_fn(
     initial_params = initial_variables['params']
     self._param_shapes = param_utils.jax_param_shapes(initial_params)
     self._param_types = param_utils.jax_param_types(self._param_shapes)
-    return jax_sharding_utils.shard(initial_params), None
+    return jax_sharding_utils.shard_along_batch_dim(initial_params), None
 
   def is_output_params(self, param_key: spec.ParameterKey) -> bool:
     return param_key == 'Dense_7'
@@ -132,11 +132,11 @@ def model_fn(
   @functools.partial(
       jax.jit,
       in_shardings=(
-          jax_sharding_utils.get_replicated_sharding(),
-          jax_sharding_utils.get_batch_sharding(),
+          jax_sharding_utils.get_replicate_sharding(),
+          jax_sharding_utils.get_batch_dim_sharding(),
       ),
       static_argnums=(0,),
-      out_shardings=jax_sharding_utils.get_replicated_sharding())
+      out_shardings=jax_sharding_utils.get_replicate_sharding())
   def _eval_batch_jitted(self,
                          params: spec.ParameterContainer,
                          batch: Dict[str, spec.Tensor]) -> spec.Tensor:

algoperf/workloads/imagenet_resnet/imagenet_jax/workload.py

@@ -103,14 +103,13 @@ def init_model_fn(
     model_state, params = pop(variables, "params")
     self._param_shapes = param_utils.jax_param_shapes(params)
     self._param_types = param_utils.jax_param_types(self._param_shapes)
-    mesh = jax_sharding_utils.get_mesh()
     params = jax.tree_map(
         lambda x: jax.device_put(x,
-                                 jax_sharding_utils.get_replicated_sharding(mesh)),
+                                 jax_sharding_utils.get_replicate_sharding()),
         params)
     model_state = jax.tree_map(
         lambda x: jax.device_put(x,
-                                 jax_sharding_utils.get_replicated_sharding(mesh)),
+                                 jax_sharding_utils.get_replicate_sharding()),
         model_state)
     return params, model_state
 
@@ -120,13 +119,13 @@ def is_output_params(self, param_key: spec.ParameterKey) -> bool:
   @functools.partial(
       jax.jit,
       in_shardings=(
-          jax_sharding_utils.get_replicated_sharding(),  # params
-          jax_sharding_utils.get_batch_sharding(),  # batch
-          jax_sharding_utils.get_replicated_sharding(),  # model_state
-          jax_sharding_utils.get_replicated_sharding(),  # rng
+          jax_sharding_utils.get_replicate_sharding(),  # params
+          jax_sharding_utils.get_batch_dim_sharding(),  # batch
+          jax_sharding_utils.get_replicate_sharding(),  # model_state
+          jax_sharding_utils.get_replicate_sharding(),  # rng
       ),
       static_argnums=(0,),
-      out_shardings=jax_sharding_utils.get_replicated_sharding())
+      out_shardings=jax_sharding_utils.get_replicate_sharding())
   def _eval_model(self,
                   params: spec.ParameterContainer,
                   batch: Dict[str, spec.Tensor],

algoperf/workloads/librispeech_conformer/librispeech_jax/workload.py

@@ -310,12 +310,12 @@ def greedy_decode(
   @functools.partial(
       jax.jit,
       in_shardings=(
-          jax_sharding_utils.get_replicated_sharding(),  # params
-          jax_sharding_utils.get_batch_sharding(),  # batch
-          jax_sharding_utils.get_replicated_sharding(),  # model_state
-          jax_sharding_utils.get_replicated_sharding(),  # rng
+          jax_sharding_utils.get_replicate_sharding(),  # params
+          jax_sharding_utils.get_batch_dim_sharding(),  # batch
+          jax_sharding_utils.get_replicate_sharding(),  # model_state
+          jax_sharding_utils.get_replicate_sharding(),  # rng
       ),
-      out_shardings=jax_sharding_utils.get_batch_sharding(),
+      out_shardings=jax_sharding_utils.get_batch_dim_sharding(),
       static_argnums=(0,))
   def _eval_step(
       self,

algoperf/workloads/librispeech_deepspeech/librispeech_jax/workload.py

@@ -108,7 +108,7 @@ def model_fn(
                                              use_running_average_bn=use_running_average_bn)
 
     model_fn_sharded = shard_map(model_fn_partial,
-                                 jax_sharding_utils.get_mesh(),
+                                 jax.sharding.Mesh(jax.devices(), ('batch')),
                                  in_specs=(None, P('batch'), None),
                                  out_specs=(P('batch'), None),
                                  )

algoperf/workloads/mnist/mnist_jax/workload.py

@@ -103,10 +103,10 @@ def loss_fn(
   @functools.partial(
       jax.jit,
       in_shardings=(
-          jax_sharding_utils.get_replicated_sharding(),  # params
-          jax_sharding_utils.get_batch_sharding(),  # batch
-          jax_sharding_utils.get_replicated_sharding(),  # model_state
-          jax_sharding_utils.get_batch_sharding(),  # rng
+          jax_sharding_utils.get_replicate_sharding(),  # params
+          jax_sharding_utils.get_batch_dim_sharding(),  # batch
+          jax_sharding_utils.get_replicate_sharding(),  # model_state
+          jax_sharding_utils.get_batch_dim_sharding(),  # rng
       ),
       static_argnums=(0,))
   def _eval_model(

algoperf/workloads/ogbg/ogbg_jax/workload.py

@@ -111,12 +111,12 @@ def _eval_metric(self, labels, logits, masks):
 
   @functools.partial(
     jax.jit,
-    in_shardings=(jax_sharding_utils.get_replicated_sharding(),
-                  jax_sharding_utils.get_batch_sharding(),
-                  jax_sharding_utils.get_replicated_sharding(),
-                  jax_sharding_utils.get_replicated_sharding()),
+    in_shardings=(jax_sharding_utils.get_replicate_sharding(),
+                  jax_sharding_utils.get_batch_dim_sharding(),
+                  jax_sharding_utils.get_replicate_sharding(),
+                  jax_sharding_utils.get_replicate_sharding()),
     static_argnums=(0,),
-    out_shardings=jax_sharding_utils.get_replicated_sharding(),
+    out_shardings=jax_sharding_utils.get_replicate_sharding(),
   )
   def _eval_batch(self, params, batch, model_state, rng):
     return super()._eval_batch(params, batch, model_state, rng)

algoperf/workloads/wmt/wmt_jax/workload.py

@@ -100,7 +100,7 @@ def eval_step(self,
   @functools.partial(
       jax.jit,
       in_shardings=(
-          jax_sharding_utils.get_batch_sharding(),  # inputs
+          jax_sharding_utils.get_batch_dim_sharding(),  # inputs
       ),
       static_argnums=(
           0,
@@ -112,9 +112,9 @@ def initialize_cache(self,
     """Initialize a cache for a given input shape and max decode length."""
     config = models.TransformerConfig(deterministic=True, decode=True)
     target_shape = (inputs.shape[0], max_decode_len) + inputs.shape[2:]
-    dummy_inputs = jax_sharding_utils.shard_naive(
+    dummy_inputs = jax_sharding_utils.shard_along_batch_dim(
         jnp.ones(inputs.shape, jnp.float32))
-    dummy_targets = jax_sharding_utils.shard_naive(
+    dummy_targets = jax_sharding_utils.shard_along_batch_dim(
         jnp.ones(target_shape, jnp.float32))
     initial_variables = models.Transformer(config).init(
         jax.random.PRNGKey(0), dummy_inputs, dummy_targets)
@@ -196,8 +196,8 @@ def translate_and_calculate_bleu(self,
         jitted_predict_step = jax.jit(
             self.predict_step,
             in_shardings=(
-                jax_sharding_utils.get_batch_sharding(),  # inputs
-                jax_sharding_utils.get_replicated_sharding(),  # params
+                jax_sharding_utils.get_batch_dim_sharding(),  # inputs
+                jax_sharding_utils.get_replicate_sharding(),  # params
                 jax_sharding_utils.get_naive_sharding_tree(cache),  # cache
             ),
             static_argnums=(
@@ -260,8 +260,8 @@ def init_model_fn(
     params_rng, dropout_rng = jax.random.split(rng)
     inputs = jnp.ones(input_shape, jnp.float32)
     targets = jnp.ones(target_shape, jnp.float32)
-    sharded_inputs = jax_sharding_utils.shard_naive(inputs)
-    sharded_targets = jax_sharding_utils.shard_naive(targets)
+    sharded_inputs = jax_sharding_utils.shard_along_batch_dim(inputs)
+    sharded_targets = jax_sharding_utils.shard_along_batch_dim(targets)
 
     initial_variables = jax.jit(
         self._eval_model.init)({'params': params_rng, 'dropout': dropout_rng},
@@ -271,7 +271,7 @@ def init_model_fn(
     initial_params = initial_variables['params']
     self._param_shapes = param_utils.jax_param_shapes(initial_params)
     self._param_types = param_utils.jax_param_types(self._param_shapes)
-    params = jax_sharding_utils.shard(initial_params)
+    params = jax_sharding_utils.shard_along_batch_dim(initial_params)
     return initial_params, None
 
   def is_output_params(self, param_key: spec.ParameterKey) -> bool: