def METHOD_NAME(model_p) -> int | None:
  """Extracts the number of cores used by the experiment.
  Args:
    model_p: The model hparams.
  Returns:
    The number of cores across all TPU slices for pjit experiments or None for
    pmap ones.
  """
  if model_p.ici_mesh_shape is None:
    return None
  def _compute_num_cores(mesh_shape) -> int:
    if mesh_shape is None:
      # Default to 1 if unset.
      return 1
    return np.prod(mesh_shape).item()
  ici_num_cores = _compute_num_cores(model_p.ici_mesh_shape)
  dcn_num_cores = _compute_num_cores(model_p.dcn_mesh_shape)
  return ici_num_cores * dcn_num_cores