code: def fill_unspecified_mesh_axes(
      parallelism_vals, target_product, parallelism_type
  ):
      """Evaluates unspecified DCN/ICI parallelism values"""
      if -1 in parallelism_vals:
          assert parallelism_vals.count(-1) == 1, (
              f"Found unspecified values (-1) for more than one {parallelism_type}   "
              "   parallelism axis. At most one axis can be unspecified."
          )
  
          determined_val = (
              target_product / np.prod(parallelism_vals) * -1
          )
  
          assert determined_val >= 1 and determined_val.is_integer, (
              "Unspecified value unable to be determined with the given     "
              f" {parallelism_type} parallelism values"
          )
  
          parallelism_vals[parallelism_vals.index(-1)] = int(
              determined_val
          )
  
      target_type = (
          "slices" if parallelism_type == "DCN" else "devices per slice"
      )
  
      assert np.prod(parallelism_vals) == target_product, (
          f"Number of {target_type} {target_product} does not match    the product"
          f" of the {parallelism_type} parallelism {np.prod(parallelism_vals)}"
      )
  
      return parallelism_vals
  
file_path: flax/examples/lm1b/utils.py
imports: import functools
  import logging
  import numpy
  import flax.linen
  from flax.linen import partitioning
  from flax.training import train_state
  import jax
  import jax.numpy
  from jax.experimental import mesh_utils
label_desc: Use image node_id 2 for calling a global function with example usage: fill_unspecified_mesh_axes(parallelism_vals, target_product, parallelism_type) and returns: parallelism_vals