//  [python]
def perturbed_alignment_score(
    sim_mat,
    gap_open,
    gap_extend,
    noise = 'normal',
    sigma = 0.1,
    num_samples = 1,
    stop_paths_gradient = True):
  """Perturbed alignment score and paths from Smith-Waterman parameters."""
  noise_sampler = functools.partial(sample_noise_with_gradients, noise)

  @tf.custom_gradient
  def forward(sw_params):
    # Perturbs the Smith-Waterman LP parameters.
    sim_mat, gap_open, gap_extend = sw_params[0], sw_params[1], sw_params[2]
    sw_params = alignment.weights_from_sim_mat(sim_mat, gap_open, gap_extend)
    shape, dtype = tf.shape(sw_params), sw_params.dtype
    pert_shape = tf.concat([[num_samples], shape], axis=0)
    noise, noise_grad = tf.nest.map_structure(
        lambda t: tf.cast(t, dtype), noise_sampler(pert_shape))
    pert_sw_params = tf.expand_dims(sw_params, 0) + sigma * noise
    pert_sw_params = tf.reshape(pert_sw_params,
                                tf.concat([[-1], shape[1:]], axis=0))
    # Computes optimal Smith-Waterman scores and alignments.
    scores, paths = hard_sw_affine(pert_sw_params)

    # Average scores and outputs over random perturbations
    def perturbation_reshape(t):
      new_shape = tf.concat([[num_samples], [-1], tf.shape(t)[1:]], 0)
      return tf.reshape(t, new_shape)

    scores = perturbation_reshape(scores)  # [num_samples, batch]
    paths = perturbation_reshape(paths)  # [num_samples, batch, len1, len2, 9]
    pert_scores = tf.reduce_mean(scores, axis=0)
    pert_paths = tf.reduce_mean(paths, axis=0)
    if stop_paths_gradient:
      pert_paths = tf.stop_gradient(pert_paths)

    def grad(ds, dp):
      # Computes grad of scores w.r.t. (packed) sw_params.
      grad_scores = pert_paths * tf.reshape(ds, [tf.shape(ds)[0], 1, 1, 1])
      # Computes grad of scores w.r.t. sim_mat, gap_open and gap_extend.
      grad_scores = alignment.adjoint_weights_from_sim_mat(
          grad_scores, gap_open.shape, gap_extend.shape)
      if stop_paths_gradient:
        gradients = grad_scores
      else:
        # Flattens paths to shape [num_samples, batch, len1 * len2 * 9].
        flat_paths = tf.reshape(
            paths, [num_samples, tf.shape(paths)[1], -1])
        flat_noise_grad = tf.reshape(
            noise_grad, [num_samples, tf.shape(noise_grad)[1], -1])
        # Flattens dp to shape [num_samples, len1 * len2 * 9]
        flat_dp = tf.reshape(dp, [tf.shape(dp)[0], -1])
        # Computes grad of paths w.r.t. (packed) sw_params.
        grad_paths = tf.einsum('nbd,nb->bd', flat_noise_grad,
                               tf.einsum('nbd,bd->nb', flat_paths, flat_dp))
        grad_paths /= sigma * num_samples
        grad_paths = tf.reshape(grad_paths, shape)  # [batch, len1, len2, 9]
        # Computes grad of paths w.r.t. sim_mat, gap_open and gap_extend.
        grad_paths = alignment.adjoint_weights_from_sim_mat(
            grad_paths, gap_open.shape, gap_extend.shape)
        # Adds gradients w.r.t. scores and gradients w.r.t. paths.
        gradients = tf.nest.map_structure(
            lambda x, y: x + y, grad_scores, grad_paths)
      return gradients

    return (pert_scores, pert_paths), grad

  return forward((sim_mat, gap_open, gap_extend))