docstr: Initialize counterfactual explanation method based on Wachter et al. (2017)
  
  Parameters
  ----------
  predict_fn
      `tensorflow` model or any other model's prediction function returning class probabilities.
  shape
      Shape of input data starting with batch size.
  distance_fn
      Distance function to use in the loss term.
  target_proba
      Target probability for the counterfactual to reach.
  target_class
      Target class for the counterfactual to reach, one of ``'other'``, ``'same'`` or an integer denoting
      desired class membership for the counterfactual instance.
  max_iter
      Maximum number of iterations to run the gradient descent for (inner loop).
  early_stop
      Number of steps after which to terminate gradient descent if all or none of found instances are solutions.
  lam_init
      Initial regularization constant for the prediction part of the Wachter loss.
  max_lam_steps
      Maximum number of times to adjust the regularization constant (outer loop) before terminating the search.
  tol
      Tolerance for the counterfactual target probability.
  learning_rate_init
      Initial learning rate for each outer loop of `lambda`.
  feature_range
      Tuple with `min` and `max` ranges to allow for perturbed instances. `Min` and `max` ranges can be `float`
      or `numpy` arrays with dimension (1 x nb of features) for feature-wise ranges.
  eps
      Gradient step sizes used in calculating numerical gradients, defaults to a single value for all
      features, but can be passed an array for feature-wise step sizes.
  init
      Initialization method for the search of counterfactuals, currently must be ``'identity'``.
  decay
      Flag to decay learning rate to zero for each outer loop over lambda.
  write_dir
      Directory to write `tensorboard` files to.
  debug
      Flag to write `tensorboard` summaries for debugging.
  sess
      Optional `tensorflow` session that will be used if passed instead of creating or inferring one internally.
docstr_tok: ["Initialize","counterfactual","explanation","method","based","on","Wachter","et","al",".","(","2017",")","Parameters","--","--","--","--","--","predict_fn","`","tensorflow","`","model","or","any","other","model","'s","prediction","function","returning","class","probabilities",".","shape","Shape","of","input","data","starting","with","batch","size",".","distance_fn","Distance","function","to","use","in","the","loss","term",".","target_proba","Target","probability","for","the","counterfactual","to","reach",".","target_class","Target","class","for","the","counterfactual","to","reach",",","one","of","``","'other","'","``",",","``","'same","'","``","or","an","integer","denoting","desired","class","membership","for","the","counterfactual","instance",".","max_iter","Maximum","number","of","iterations","to","run","the","gradient","descent","for","(","inner","loop",")",".","early_stop","Number","of","steps","after","which","to","terminate","gradient","descent","if","all","or","none","of","found","instances","are","solutions",".","lam_init","Initial","regularization","constant","for","the","prediction","part","of","the","Wachter","loss",".","max_lam_steps","Maximum","number","of","times","to","adjust","the","regularization","constant","(","outer","loop",")","before","terminating","the","search",".","tol","Tolerance","for","the","counterfactual","target","probability",".","learning_rate_init","Initial","learning","rate","for","each","outer","loop","of","`","lambda","`",".","feature_range","Tuple","with","`","min","`","and","`","max","`","ranges","to","allow","for","perturbed","instances",".","`","Min","`","and","`","max","`","ranges","can","be","`","float","`","or","`","numpy","`","arrays","with","dimension","(","1","x","nb","of","features",")","for","feature-wise","ranges",".","eps","Gradient","step","sizes","used","in","calculating","numerical","gradients",",","defaults","to","a","single","value","for","all","features",",","but","can","be","passed","an","array","for","feature-wise","step","sizes",".","init","Initialization","method","for","the","search","of","counterfactuals",",","currently","must","be","``","'identity","'","``",".","decay","Flag","to","decay","learning","rate","to","zero","for","each","outer","loop","over","lambda",".","write_dir","Directory","to","write","`","tensorboard","`","files","to",".","debug","Flag","to","write","`","tensorboard","`","summaries","for","debugging",".","sess","Optional","`","tensorflow","`","session","that","will","be","used","if","passed","instead","of","creating","or","inferring","one","internally","."]
code: def __init__(
      self,
      predict_fn: Union[
          Callable[[np.ndarray], np.ndarray], tf.keras.Model
      ],
      shape: Tuple[int, ...],
      distance_fn: str = "l1",
      target_proba: float = 1.0,
      target_class: Union[str, int] = "other",
      max_iter: int = 1000,
      early_stop: int = 50,
      lam_init: float = 1e-1,
      max_lam_steps: int = 10,
      tol: float = 0.05,
      learning_rate_init=0.1,
      feature_range: Union[Tuple, str] = (-1e10, 1e10),
      eps: Union[float, np.ndarray] = 0.01,  # feature-wise epsilons
      init: str = "identity",
      decay: bool = True,
      write_dir: Optional[str] = None,
      debug: bool = False,
      sess: Optional[tf.Session] = None,
  ) -> None:
      """
      Initialize counterfactual explanation method based on Wachter et al. (2017)
  
      Parameters
      ----------
      predict_fn
          `tensorflow` model or any other model's prediction function returning class probabilities.
      shape
          Shape of input data starting with batch size.
      distance_fn
          Distance function to use in the loss term.
      target_proba
          Target probability for the counterfactual to reach.
      target_class
          Target class for the counterfactual to reach, one of ``'other'``, ``'same'`` or an integer denoting
          desired class membership for the counterfactual instance.
      max_iter
          Maximum number of iterations to run the gradient descent for (inner loop).
      early_stop
          Number of steps after which to terminate gradient descent if all or none of found instances are solutions.
      lam_init
          Initial regularization constant for the prediction part of the Wachter loss.
      max_lam_steps
          Maximum number of times to adjust the regularization constant (outer loop) before terminating the search.
      tol
          Tolerance for the counterfactual target probability.
      learning_rate_init
          Initial learning rate for each outer loop of `lambda`.
      feature_range
          Tuple with `min` and `max` ranges to allow for perturbed instances. `Min` and `max` ranges can be `float`
          or `numpy` arrays with dimension (1 x nb of features) for feature-wise ranges.
      eps
          Gradient step sizes used in calculating numerical gradients, defaults to a single value for all
          features, but can be passed an array for feature-wise step sizes.
      init
          Initialization method for the search of counterfactuals, currently must be ``'identity'``.
      decay
          Flag to decay learning rate to zero for each outer loop over lambda.
      write_dir
          Directory to write `tensorboard` files to.
      debug
          Flag to write `tensorboard` summaries for debugging.
      sess
          Optional `tensorflow` session that will be used if passed instead of creating or inferring one internally.
      """
      super().__init__(meta=copy.deepcopy(DEFAULT_META_CF))
      # get params for storage in meta
      params = locals()
      remove = ["self", "predict_fn", "sess", "__class__"]
      for key in remove:
          params.pop(key)
      self.meta["params"].update(params)
  
      self.data_shape = shape
      self.batch_size = shape[0]
      self.target_class = target_class
  
      # options for the optimizer
      self.max_iter = max_iter
      self.lam_init = lam_init
      self.tol = tol
      self.max_lam_steps = max_lam_steps
      self.early_stop = early_stop
  
      self.eps = eps
      self.init = init
      self.feature_range = feature_range
      self.target_proba_arr = target_proba * np.ones(self.batch_size)
  
      self.debug = debug
  
      # check if the passed object is a model and get session
      is_model = isinstance(predict_fn, tf.keras.Model)
      model_sess = tf.compat.v1.keras.backend.get_session()
  
      self.meta["params"].update(is_model=is_model)
  
      # if session provided, use it
      if isinstance(sess, tf.Session):
          self.sess = sess
      else:
          self.sess = model_sess
  
      if is_model:  # Keras or TF model
          self.model = True
          self.predict_fn = predict_fn.predict  # type: ignore # array function
          self.predict_tn = predict_fn  # tensor function
  
      else:  # black-box model
          self.predict_fn = predict_fn
          self.predict_tn = None
          self.model = False
  
      self.n_classes = self.predict_fn(np.zeros(shape)).shape[1]
  
      # flag to keep track if explainer is fit or not
      self.fitted = False
  
      # set up graph session for optimization (counterfactual search)
      with tf.variable_scope("cf_search", reuse=tf.AUTO_REUSE):
          # define variables for original and candidate counterfactual instances, target labels and lambda
          self.orig = tf.get_variable(
              "original", shape=shape, dtype=tf.float32
          )
          self.cf = tf.get_variable(
              "counterfactual",
              shape=shape,
              dtype=tf.float32,
              constraint=lambda x: tf.clip_by_value(
                  x, feature_range[0], feature_range[1]
              ),
          )
          # the following will be a 1-hot encoding of the target class (as predicted by the model)
          self.target = tf.get_variable(
              "target",
              shape=(self.batch_size, self.n_classes),
              dtype=tf.float32,
          )
  
          # constant target probability and global step variable
          self.target_proba = tf.constant(
              target_proba * np.ones(self.batch_size),
              dtype=tf.float32,
              name="target_proba",
          )
          self.global_step = tf.Variable(
              0.0, trainable=False, name="global_step"
          )
  
          # lambda hyperparameter - placeholder instead of variable as annealed in first epoch
          self.lam = tf.placeholder(
              tf.float32, shape=(self.batch_size), name="lam"
          )
  
          # define placeholders that will be assigned to relevant variables
          self.assign_orig = tf.placeholder(
              tf.float32, shape, name="assing_orig"
          )
          self.assign_cf = tf.placeholder(
              tf.float32, shape, name="assign_cf"
          )
          self.assign_target = tf.placeholder(
              tf.float32,
              shape=(self.batch_size, self.n_classes),
              name="assign_target",
          )
  
          # L1 distance and MAD constants
          # TODO: MADs?
          ax_sum = list(np.arange(1, len(self.data_shape)))
          if distance_fn == "l1":
              self.dist = tf.reduce_sum(
                  tf.abs(self.cf - self.orig), axis=ax_sum, name="l1"
              )
          else:
              logger.exception(
                  "Distance metric %s not supported", distance_fn
              )
              raise ValueError
  
          # distance loss
          self.loss_dist = self.lam * self.dist
  
          # prediction loss
          if not self.model:
              # will need to calculate gradients numerically
              self.loss_opt = self.loss_dist
          else:
              # autograd gradients throughout
              self.pred_proba = self.predict_tn(self.cf)
  
              # 3 cases for target_class
              if target_class == "same":
                  self.pred_proba_class = tf.reduce_max(
                      self.target * self.pred_proba, 1
                  )
              elif target_class == "other":
                  self.pred_proba_class = tf.reduce_max(
                      (1 - self.target) * self.pred_proba, 1
                  )
              elif target_class in range(self.n_classes):
                  # if class is specified, this is known in advance
                  self.pred_proba_class = tf.reduce_max(
                      tf.one_hot(
                          target_class, self.n_classes, dtype=tf.float32
                      )
                      * self.pred_proba,
                      1,
                  )
              else:
                  logger.exception(
                      "Target class %s unknown", target_class
                  )
                  raise ValueError
  
              self.loss_pred = tf.square(
                  self.pred_proba_class - self.target_proba
              )
  
              self.loss_opt = self.loss_pred + self.loss_dist
  
          # optimizer
          if decay:
              self.learning_rate = tf.train.polynomial_decay(
                  learning_rate_init,
                  self.global_step,
                  self.max_iter,
                  0.0,
                  power=1.0,
              )
          else:
              self.learning_rate = tf.convert_to_tensor(
                  learning_rate_init
              )
  
          # TODO optional argument to change type, learning rate scheduler
          opt = tf.train.AdamOptimizer(self.learning_rate)
  
          # first compute gradients, then apply them
          self.compute_grads = opt.compute_gradients(
              self.loss_opt, var_list=[self.cf]
          )
          self.grad_ph = tf.placeholder(
              shape=shape, dtype=tf.float32, name="grad_cf"
          )
          grad_and_var = [(self.grad_ph, self.cf)]
          self.apply_grads = opt.apply_gradients(
              grad_and_var, global_step=self.global_step
          )
  
      # variables to initialize
      self.setup: list = []
      self.setup.append(self.orig.assign(self.assign_orig))
      self.setup.append(self.cf.assign(self.assign_cf))
      self.setup.append(self.target.assign(self.assign_target))
  
      self.tf_init = tf.variables_initializer(
          var_list=tf.global_variables(scope="cf_search")
      )
  
      # tensorboard
      if write_dir is not None:
          self.writer = tf.summary.FileWriter(
              write_dir, tf.get_default_graph()
          )
          self.writer.add_graph(tf.get_default_graph())
  
      # return templates
      self.instance_dict = dict.fromkeys(
          ["X", "distance", "lambda", "index", "class", "proba", "loss"]
      )
      self.return_dict = copy.deepcopy(DEFAULT_DATA_CF)
      self.return_dict["all"] = {
          i: [] for i in range(self.max_lam_steps)
      }
  
imports: import copy
  import logging
  from typing import Callable, Optional, Tuple, Union
  import numpy
  import tensorflow.compat.v1
  from alibi.api.defaults import DEFAULT_DATA_CF, DEFAULT_META_CF
  from alibi.api.interfaces import Explainer, Explanation
  from alibi.utils.gradients import num_grad_batch
label_desc: Use image node_id 1 to create a new Counterfactual object from inherited base classes: Explainer with example: obj = Counterfactual(predict_fn, shape, distance_fn, target_proba, target_class, max_iter, early_stop, lam_init, max_lam_steps, tol, learning_rate_init, feature_range, eps, init, decay, write_dir, debug, sess)