(
    objective_fn,
    module_variables_init,
    module_variables_final,
    num_samples_per_iteration=10,
    alpha_grid_size=10,
    sigma_grid_size=10,
    sigma_ratio=1.0,
    loss_threshold_condition=relative_error_condition,
    normalize_error=False,
)
return ModuleCriticalityAnalysis(
      criticality_score=norm_final,
      alpha=alpha_final,
      sigma=sigma_final,
      loss_value=loss_final,
      num_samples_per_iteration=num_samples_per_iteration,
      alpha_grid_size=alpha_grid_size,
      sigma_grid_size=sigma_grid_size,
      sigma_ratio=sigma_ratio,
      initial_objective_value=initial_objective_value,
      final_objective_value=final_objective_value,
  )
Compute the criticality of a module parameterized by `module_variables`.

  Args:
    objective_fn: A callable that takes in an iterable of the module-specific
      variables and produces the value of the objective function.
    module_variables_init: A list of tf.Tensors; the variables of the module at
      initialization.
    module_variables_final: A list of tf.Tensors; the variables of the module at
      convergence.
    num_samples_per_iteration: Number of perturbations to sample each iteration.
    alpha_grid_size: The number of values to test for alpha, the interpolation
      coefficient.
    sigma_grid_size: The number of values to test for sigma, the standard
      deviation of the perturbation.
    sigma_ratio: Positive scalar multiplier k for values of sigma, to enforce
      that the tested values of sigma lie in [k * 1e-16, k]; the default is 1.0,
      implying that the tested values of sigma lie in the interval [1e-16, 1].
    loss_threshold_condition: A callable that takes in a reference objective
      value and a candidate objective value and produces a thresholding
      decision.
    normalize_error: Whether to normalize the error that is minimized over in
      the definition of criticality by the Frobenius norm of the distance
      between initial and final parameters.

  Returns:
    A `collections.NamedTuple` that contains the results of the criticality
    analysis.
Compute the criticality of a module parameterized by `module_variables`.
Compute the criticality of a module parameterized by module_variables .
def compute_module_criticality(
    objective_fn,
    module_variables_init,
    module_variables_final,
    num_samples_per_iteration=10,
    alpha_grid_size=10,
    sigma_grid_size=10,
    sigma_ratio=1.0,
    loss_threshold_condition=relative_error_condition,
    normalize_error=False,
):
  """Compute the criticality of a module parameterized by `module_variables`.

  Args:
    objective_fn: A callable that takes in an iterable of the module-specific
      variables and produces the value of the objective function.
    module_variables_init: A list of tf.Tensors; the variables of the module at
      initialization.
    module_variables_final: A list of tf.Tensors; the variables of the module at
      convergence.
    num_samples_per_iteration: Number of perturbations to sample each iteration.
    alpha_grid_size: The number of values to test for alpha, the interpolation
      coefficient.
    sigma_grid_size: The number of values to test for sigma, the standard
      deviation of the perturbation.
    sigma_ratio: Positive scalar multiplier k for values of sigma, to enforce
      that the tested values of sigma lie in [k * 1e-16, k]; the default is 1.0,
      implying that the tested values of sigma lie in the interval [1e-16, 1].
    loss_threshold_condition: A callable that takes in a reference objective
      value and a candidate objective value and produces a thresholding
      decision.
    normalize_error: Whether to normalize the error that is minimized over in
      the definition of criticality by the Frobenius norm of the distance
      between initial and final parameters.

  Returns:
    A `collections.NamedTuple` that contains the results of the criticality
    analysis.
  """
  initial_objective_value = objective_fn(module_variables_init)
  final_objective_value = objective_fn(module_variables_final)

  # Test a 2D grid of alpha and sigma values.
  float_zero = tf.cast(0, tf.float32)
  alphas, sigmas = tf.meshgrid(
      tf.linspace(float_zero, 1, alpha_grid_size + 1),
      tf.linspace(float_zero + 1e-16, 1, sigma_grid_size + 1) * sigma_ratio,
  )
  alphas, sigmas = tf.reshape(alphas, [-1]), tf.reshape(sigmas, [-1])

  def _evaluate_alpha_sigma(alpha_sigma):
    alpha, sigma = alpha_sigma
    return _interpolate_and_perturb(
        alpha=alpha,
        sigma=sigma,
        params_init=module_variables_init,
        params_final=module_variables_final,
        objective_fn=objective_fn,
        loss_threshold_condition=functools.partial(
            loss_threshold_condition, reference_error=final_objective_value),
        normalize_error=normalize_error,
        num_samples_per_iteration=num_samples_per_iteration,
    )

  (threshold_conditions, interpolated_and_perturbed_losses,
   interpolated_and_perturbed_norms) = tf.map_fn(
       _evaluate_alpha_sigma,
       elems=(alphas, sigmas),
       dtype=(tf.bool, tf.float32, tf.float32),
   )

  masked_interpolated_and_perturbed_norms = tf.where(
      threshold_conditions, interpolated_and_perturbed_norms,
      tf.ones_like(interpolated_and_perturbed_norms) * np.inf)
  idx_min = tf.math.argmin(masked_interpolated_and_perturbed_norms)
  (loss_final, norm_final, alpha_final,
   sigma_final) = (interpolated_and_perturbed_losses[idx_min],
                   interpolated_and_perturbed_norms[idx_min], alphas[idx_min],
                   sigmas[idx_min])

  return ModuleCriticalityAnalysis(
      criticality_score=norm_final,
      alpha=alpha_final,
      sigma=sigma_final,
      loss_value=loss_final,
      num_samples_per_iteration=num_samples_per_iteration,
      alpha_grid_size=alpha_grid_size,
      sigma_grid_size=sigma_grid_size,
      sigma_ratio=sigma_ratio,
      initial_objective_value=initial_objective_value,
      final_objective_value=final_objective_value,
  )
def compute_module_criticality ( objective_fn , module_variables_init , module_variables_final , num_samples_per_iteration = 10 , alpha_grid_size = 10 , sigma_grid_size = 10 , sigma_ratio = 1.0 , loss_threshold_condition = relative_error_condition , normalize_error = False , ) : initial_objective_value = objective_fn ( module_variables_init ) final_objective_value = objective_fn ( module_variables_final ) # Test a 2D grid of alpha and sigma values. float_zero = tf . cast ( 0 , tf . float32 ) alphas , sigmas = tf . meshgrid ( tf . linspace ( float_zero , 1 , alpha_grid_size + 1 ) , tf . linspace ( float_zero + 1e-16 , 1 , sigma_grid_size + 1 ) * sigma_ratio , ) alphas , sigmas = tf . reshape ( alphas , [ - 1 ] ) , tf . reshape ( sigmas , [ - 1 ] ) def _evaluate_alpha_sigma ( alpha_sigma ) : alpha , sigma = alpha_sigma return _interpolate_and_perturb ( alpha = alpha , sigma = sigma , params_init = module_variables_init , params_final = module_variables_final , objective_fn = objective_fn , loss_threshold_condition = functools . partial ( loss_threshold_condition , reference_error = final_objective_value ) , normalize_error = normalize_error , num_samples_per_iteration = num_samples_per_iteration , ) ( threshold_conditions , interpolated_and_perturbed_losses , interpolated_and_perturbed_norms ) = tf . map_fn ( _evaluate_alpha_sigma , elems = ( alphas , sigmas ) , dtype = ( tf . bool , tf . float32 , tf . float32 ) , ) masked_interpolated_and_perturbed_norms = tf . where ( threshold_conditions , interpolated_and_perturbed_norms , tf . ones_like ( interpolated_and_perturbed_norms ) * np . inf ) idx_min = tf . math . argmin ( masked_interpolated_and_perturbed_norms ) ( loss_final , norm_final , alpha_final , sigma_final ) = ( interpolated_and_perturbed_losses [ idx_min ] , interpolated_and_perturbed_norms [ idx_min ] , alphas [ idx_min ] , sigmas [ idx_min ] ) return ModuleCriticalityAnalysis ( criticality_score = norm_final , alpha = alpha_final , sigma = sigma_final , loss_value = loss_final , num_samples_per_iteration = num_samples_per_iteration , alpha_grid_size = alpha_grid_size , sigma_grid_size = sigma_grid_size , sigma_ratio = sigma_ratio , initial_objective_value = initial_objective_value , final_objective_value = final_objective_value , )