<filename>binary_logreg.py
from cr_model import CRModel
import warnings

import jax.numpy as jnp
import jax
from jax.numpy.linalg import norm
from jax import jit, value_and_grad, random, custom_jvp
from jax.experimental import stax
from jax.experimental.stax import Dense
from jax.nn import log_sigmoid, sigmoid
from jax.nn.initializers import normal, zeros
from jax.flatten_util import ravel_pytree

from tensorflow_probability.substrates.jax.optimizer import lbfgs_minimize

from util import register_pytree_node_dataclass, tree_zeros_like
from dataclasses import dataclass
import dataclasses
from typing import Dict, Union, Any, Optional, Tuple

from functools import partial

Dataset = Tuple[jnp.ndarray, jnp.ndarray]
def _zeros_64(key, shape): return zeros(key, shape, dtype=jnp.dtype('float64'))


_init_params, _decision_function = stax.serial(
            Dense(1, W_init=_zeros_64, b_init=_zeros_64)
        )


def _init_random_params(rng: jnp.ndarray, shape: Any, sigma: float):
    """Initializer for randomized coefficients of the linear perturbation term"""
    params = []
    r1, r2 = random.split(rng)
    norm_init = normal(stddev=sigma, dtype=jnp.dtype('float64'))
    coef = norm_init(r1, (shape[-1], 1))
    intercept = norm_init(r2, (1,))
    param = coef, intercept
    params.append(param)
    return shape, params


def _init_data_weights(data: Dataset, data_weights: Optional[jnp.ndarray]) -> jnp.ndarray:
    """Helper function that returns initialized sample weights. If not specified, they are initialized to one."""
    if data_weights is None:
        return jnp.ones(data[0].shape[0], dtype=data[0].dtype)
    else:
        return data_weights


@partial(jit, static_argnums=(2,))
def _init_gram_matrix(inputs: jnp.ndarray, weights: jnp.ndarray, add_intercept: bool = True):
    """Compute the gram matrix for the inputs

    Args:
        inputs: 2d array of shape (n_samples, n_features) which contains the input features.
        weights: 1d array of shape (n_samples,) which contains non-negative weights for each sample.
        add_intercept: whether to append an extra constant feature to incorporate an intercept. Defaults to True.

    Returns:
        a 2d array of shape (n_features, n_features) if add_intercept is False, otherwise a 2d array of shape
        (n_features + 1, n_features + 1).
    """
    inputs_weighted = inputs * weights[:, None]
    gram_matrix = inputs_weighted.T @ inputs_weighted

    # Account for intercept
    if add_intercept:
        inputs_colsums = jnp.sum(inputs_weighted, axis=0)
        gram_matrix = jnp.block([[gram_matrix, inputs_colsums.reshape(-1, 1)],
                                 [inputs_colsums, jnp.sum(weights)]])

    return gram_matrix


@jit
def _objective(params: Any, random_params: Any, data: Dataset, lamb: float, pos_label: Union[str, int, float],
               data_weights: Optional[jnp.ndarray]) -> float:
    """Objective function for L2-regularized logistic regression. Includes a randomized perturbation term to
    improve privacy."""
    data_weights = _init_data_weights(data, data_weights)

    inputs = jnp.atleast_2d(data[0])
    targets = jnp.atleast_1d(data[1])
    targets = targets.ravel()  # ensure targets is a 1d vector
    scores = _decision_function(params, inputs).ravel()
    cross_entropy_terms = jnp.where(targets == pos_label, log_sigmoid(scores), log_sigmoid(-scores))
    cross_entropy_loss = -jnp.dot(data_weights, cross_entropy_terms)
    
    # Regularization penalty term
    coef, intercept = params[0]
    # Guo et al. (2020) penalize the coefficients *and* the bias
    l2_penalty = 0.5 * lamb * jnp.sum(data_weights) * (jnp.sum(coef**2) + jnp.sum(intercept**2))

    # Random linear term for indistinguishability/privacy
    b_coef, b_intercept = random_params[0]
    # Note: this term is divided by data_size in
    # https://github.com/facebookresearch/certified-removal/blob/2a3aa66e85b95d659830f944f644daaf41f63167/test_removal.py#L109
    rand_linear = (b_coef * coef).sum() + (b_intercept * intercept).sum()
    
    return cross_entropy_loss + l2_penalty + rand_linear


@partial(jit, static_argnums=3)
def _gradient(model: 'BinaryLogReg', data: Dataset, data_weights: Optional[jnp.ndarray],
              perturbation_term: bool) -> jnp.ndarray:
    """Ravelled gradient of the objective function with respect to the model parameters"""
    params_flat, unravel = ravel_pytree(model.params)
    random_params = model.random_params
    if not perturbation_term:
        random_params = tree_zeros_like(random_params)
    g = jax.grad(lambda p: _objective(unravel(p), random_params, data, model.lamb, model.pos_label, data_weights))
    return g(params_flat)


@jit
def _hessian(model: 'BinaryLogReg', data: Dataset, data_weights: Optional[jnp.ndarray]) -> jnp.ndarray:
    """Ravelled Hessian matrix of the objective function with respect to the model parameters"""
    params_flat, unravel = ravel_pytree(model.params)
    random_params = model.random_params
    h = jax.hessian(lambda p: _objective(unravel(p), random_params, data, model.lamb, model.pos_label, data_weights))
    return h(params_flat)


@custom_jvp
def _fit(init_params: Any, random_params: Any, data: Dataset, lamb: float, pos_label: Union[str, int, float],
         data_weights: jnp.ndarray, tolerance: float, max_iterations: int) -> Tuple[Any, Dict[str, Any]]:
    """Implementation of the `fit` method for BinaryLogReg"""
    # L-BFGS implementation in TFP is not compatible with PyTrees, so need to 
    # ravel parameters. Call ravelled parameters "position".
    init_position, unravel = ravel_pytree(init_params)
    
    # Wrapper for objective function which accepts position as argument.
    def obj(position):
        params = unravel(position)
        return _objective(params, random_params, data, lamb, pos_label, data_weights)
    
    obj_and_grad = value_and_grad(obj)

    def solve(f, init_position):
        return lbfgs_minimize(f, init_position, tolerance=tolerance, max_iterations=max_iterations)
    
    result = solve(obj_and_grad, init_position)
    params = unravel(result.position)

    return params, {'converged': result.converged, 
                    'num_iterations': result.num_iterations, 
                    'grad_norm': jnp.linalg.norm(result.objective_gradient, ord=2)}


# Custom implementation of jvp based on implicit function theorem
@_fit.defjvp
def _fit_jvp(primals, tangents):
    init_params, random_params, data, lamb, pos_label, data_weights, tolerance, max_iterations = primals
    _, _, data_dot, _, _, _, _, _ = tangents
    
    params, diagnostics = _fit(init_params, random_params, data, lamb, pos_label, data_weights, tolerance,
                               max_iterations)
    
    inputs, targets = data
    inputs_dot, _ = data_dot

    scores = _decision_function(params, inputs).ravel()
    preds = sigmoid(scores)
    preds_var = preds * (1 - preds)
    scores_dot = _decision_function(params, inputs_dot).ravel()
    
    first_term = jnp.einsum('i,i,ij->j', scores_dot, preds_var, inputs)
    first_term = jnp.concatenate((jnp.atleast_1d(inputs.shape[0]), first_term))
    second_term = jnp.einsum('i,ij->j', preds - targets, inputs_dot)
    second_term = jnp.concatenate((jnp.atleast_1d(0), second_term))
    jac_X_grad = first_term + second_term
    
    params_flat, unravel = ravel_pytree(params)
    hess = jax.hessian(lambda p: _objective(unravel(p), random_params, data, lamb, pos_label, data_weights))(params_flat)
    hess_inv = jnp.linalg.inv(hess)
    
    # Pre-multiplying by the inverse Jacobian of the gradient 
    # of the loss w.r.t. w (i.e. the inverse Hessian) gives 
    # the required result 
    tangent_out = unravel(hess_inv @ jac_X_grad), {k: jnp.zeros_like(jnp.asarray(v)) for k, v in diagnostics.items()}
    
    primal_out = params, diagnostics
    
    return primal_out, tangent_out


_fit_jit = jax.jit(_fit)


@partial(jit, static_argnums=5)
def _unlearn(model: 'BinaryLogReg', data: Dataset, delete: jnp.ndarray, retain: jnp.ndarray,
             prev_gram_matrix: Optional[jnp.ndarray], use_full_data_hess_approx: bool) -> \
        Tuple[jnp.ndarray, float, jnp.ndarray]:
    """Implementation of the `unlearn` method for BinaryLogReg"""
    inputs, _ = data
    
    # Gradient contribution due to deleted data
    grad_del = BinaryLogReg.gradient(model, data, delete, False)
    
    # Correct error in the quadratic approximation used by Guo et al. (2020). Namely, the Hessian should be
    # approximated on the objective prior to removal.
    if use_full_data_hess_approx:
        retain = jnp.logical_or(delete, retain)
    
    # Update the Gram matrix.
    if prev_gram_matrix is None:
        # Compute for the first time
        gram_matrix = _init_gram_matrix(inputs, retain, add_intercept=True)
    else:
        gram_matrix = prev_gram_matrix
        if not use_full_data_hess_approx:
            delta_gram_matrix = _init_gram_matrix(inputs, delete, add_intercept=True)
            gram_matrix = gram_matrix - delta_gram_matrix
    
    # Although the Hessian can be expressed as a sum over training instances, we can't update it incrementally. This
    # is because the model parameters change at each iteration.
    hess = BinaryLogReg.hessian(model, data, retain)

    hess_inv_grad_del = jnp.linalg.inv(hess) @ grad_del
    inputs_x_hess_inv_grad_del = (inputs * retain[:, None]) @ hess_inv_grad_del[0:-1] + hess_inv_grad_del[-1]
    delta_grnb = (0.25 * jnp.sqrt(norm(gram_matrix, ord=2)) * norm(hess_inv_grad_del, ord=2) *
                  norm(inputs_x_hess_inv_grad_del, ord=2))
    
    # Easier to perform a Newton update using ravelled parameters
    params_flat, unravel = ravel_pytree(model.params)
    params = unravel(params_flat + hess_inv_grad_del)
    
    # Update GRNB
    grnb = model.grnb + delta_grnb

    return params, grnb, gram_matrix


@register_pytree_node_dataclass
@dataclass(frozen=True)
class BinaryLogReg(CRModel):
    """Logistic regression model with support for certified removal

    Attributes:
        lamb: L2 regularization strength.
        params: model parameters, represented as a tuple containing a vector of coefficients and an intercept.
        random_params: random parameters, represented as a tuple containing a vector of coefficients and an
            intercept.
        epsilon: parameter for the certified removal guarantee, typically on the unit interval.
        delta: parameter for the certified removal guarantee, typically of order 1/(training set size).
        sigma: standard deviation of the random parameters.
        pos_label: label used to encode positive instances.
        neg_label: label used to encode negative instances.
        grnb: cumulative gradient residual norm bound.
        grnb_thres: threshold on `grnb` above which full retraining is triggered.
    """
    lamb: float = 1.0
    params: Optional[Tuple[jnp.ndarray, jnp.ndarray]] = None
    random_params: Optional[Tuple[jnp.ndarray, jnp.ndarray]] = None
    epsilon: float = 1.0
    delta: float = 1e-4
    sigma: float = 1.0
    pos_label: Union[str, int, float] = 1
    neg_label: Union[str, int, float] = 0
    grnb: float = 0.0
    
    @property
    def grnb_thres(self) -> float:
        return self.sigma * self.epsilon / jnp.sqrt(2 * jnp.log(1.5 / self.delta))

    def decision_function(self, inputs: jnp.ndarray) -> jnp.ndarray:
        """Predict confidence scores

        The confidence score for a sample is proportional to the signed distance of that sample to the hyperplane.

        Args:
            inputs: sample inputs to score. Must be a 2d array of shape (n_samples, n_features).

        Returns:
            array of shape (n_samples,) containing the confidence scores (for the positive class) for each sample.
            A positive score means the positive class would be predicted.
        """
        assert self.params is not None, "model is not fitted"
        return _decision_function(self.params, inputs)

    def predict(self, inputs: jnp.ndarray) -> jnp.ndarray:
        scores = self.decision_function(inputs)
        binary = (scores > 0)
        return jnp.where(binary, self.pos_label, self.neg_label)

    def predict_proba(self, inputs: jnp.ndarray) -> jnp.ndarray:
        """Probability estimates

        Args:
            inputs: samples inputs on which to estimate predicted probabilities. Must be a 2d array of shape
                (n_samples, n_features).

        Returns:
            array of shape (n_samples,) containing the predicted probabilities that each sample is in the positive
            class.
        """
        scores = self.decision_function(inputs)
        return sigmoid(scores)

    def fit(self, data: Dataset, rng: jnp.ndarray, data_weights: Optional[jnp.ndarray] = None,
            tolerance: float = 1e-4, max_iterations: int = 1000) -> Tuple['BinaryLogReg', Dict]:
        """Fit the model

        Args:
            data: training data represented as a tuple, where the first element is the inputs and the second element
                is the targets (labels).
            rng: a PRNGKey (an array with shape (2,) and dtype uint32).
            data_weights: non-negative weights associated with the samples in `data`.
            tolerance: gradient tolerance for L-BFGS. If the sup-norm of the gradient vector is below this value,
                L-BFGS is terminated.
            max_iterations: maximum number of iterations before L-BFGS is terminated.

        Returns:
            a tuple where the first element is the fitted model, and the second element contains diagnostics
            associated with L-BFGS.
        """
        data_weights = _init_data_weights(data, data_weights)

        inputs, _ = data
        
        # Update random parameters
        rng1, rng2 = random.split(rng)
        _, init_params = _init_params(rng1, (-1, inputs.shape[1]))
        _, random_params = _init_random_params(rng2, (-1, inputs.shape[1]), self.sigma)

        params, diagnostics = _fit_jit(init_params, random_params, data, self.lamb, self.pos_label,
                                       data_weights, tolerance, max_iterations)

        # Update model
        model_dict = dataclasses.asdict(self)
        model_dict['params'] = params
        model_dict['random_params'] = random_params
        model_dict['grnb'] = 0.0
        model = BinaryLogReg(**model_dict)

        return model, diagnostics

    def unlearn(self, data: Dataset, delete: jnp.ndarray, retain: jnp.ndarray, rng: jnp.ndarray,
                tolerance: float = 1e-4, max_iterations: int = 100, prev_gram_matrix: Optional[jnp.ndarray] = None,
                enforce_grnb_constraint: bool = True, use_full_data_hess_approx: bool = True,
                warn_retrain: bool = False) -> Tuple['BinaryLogReg', Optional[Dict], Optional[jnp.ndarray], bool]:
        """Unlearn training data

        Args:
            data: training data represented as a tuple, where the first element is the inputs and the second element
                is the targets (labels).
            delete: boolean array specifying samples in data to unlearn/delete.
            retain: boolean array specifying samples in data to keep post-deletion.
            rng: a PRNGKey (an array with shape (2,) and dtype uint32).
            tolerance: gradient tolerance for L-BFGS. If the sup-norm of the gradient vector is below this value,
                L-BFGS is terminated.
            max_iterations: maximum number of iterations before L-BFGS is terminated.
            prev_gram_matrix: gram matrix for the training data inputs prior to unlearning. If not provided, it is
                computed from scratch.
            enforce_grnb_constraint: boolean specifying whether to retrain from scratch if the cumulative gradient
                residual norm exceeds the bound that guarantees (epsilon, delta)-CR.
            use_full_data_hess_approx: boolean specifying whether to use the full training data (prior to deletion)
                to approximate the Hessian. If False, only the retained data (post-deletion) is used to approximate
                the Hessian as in Guo et al. (2020).
            warn_retrain: boolean specifying whether to emit a warning if retraining from scratch is triggered.
            
        Returns:
            a tuple where the first element is the unlearned model, the second element contains diagnostics associated
            with L-BFGS for retraining, the third element is the updated gram matrix, and the fourth element
            is a boolean indicating whether retraining was triggered.
        """
        assert self.params is not None, "model is not fitted"
        params, grnb, gram_matrix = _unlearn(self, data, delete, retain, prev_gram_matrix, use_full_data_hess_approx)
        
        def retrain(_):
            if warn_retrain:
                warnings.warn("GRNB of {} exceeds upper bound of {}. Retraining from "
                              "scratch".format(grnb, self.grnb_thres))
            model_new, diagnostics = BinaryLogReg.fit(self, data, rng, data_weights=retain, tolerance=tolerance,
                                                      max_iterations=max_iterations)
            return model_new, diagnostics, gram_matrix, True
    
        def done(_):
            # Update model
            model_dict = dataclasses.asdict(self)
            model_dict['params'] = params
            model_dict['grnb'] = grnb
            model_new = BinaryLogReg(**model_dict)
            diagnostics = {'converged': jnp.zeros((), dtype=bool), 'num_iterations': jnp.zeros((), dtype=jnp.int32),
                           'grad_norm': jnp.zeros((), dtype=jnp.float64)}
            return model_new, diagnostics, gram_matrix, False
        
        return jax.lax.cond(
            jnp.any(grnb > self.grnb_thres) & enforce_grnb_constraint,
            retrain,
            done,
            None
        )
    
    def set_sigma(self, value: float) -> 'BinaryLogReg':
        """Set sigma parameter
        
        Args:
            value: new value of sigma
        
        Returns:
            new model with sigma parameter updated
        """
        return dataclasses.replace(self, sigma=value, params=None, random_params=None, grnb=0.0)

    def objective(self, data: Dataset, data_weights: Optional[jnp.ndarray] = None) -> float:
        """Compute the objective on the given training data"""
        assert self.params is not None, "model is not fitted"
        return _objective(self.params, self.random_params, data, self.lamb, self.pos_label, data_weights)

    def gradient(self, data: Dataset, data_weights: Optional[jnp.ndarray] = None,
                 perturbation_term: bool = True) -> jnp.ndarray:
        """Compute the (ravelled) gradient of the objective with respect to the model parameters on the given
        training data"""
        assert self.params is not None, "model is not fitted"
        return _gradient(self, data, data_weights, perturbation_term)

    def hessian(self, data: Dataset, data_weights: Optional[jnp.ndarray] = None) -> jnp.ndarray:
        """Compute the (ravelled) Hessian of the objective with respect to the model parameters on the given training
        data"""
        assert self.params is not None, "model is not fitted"
        return _hessian(self, data, data_weights)

<filename>binary_logreg.py
from cr_model import CRModel
import warnings

import jax.numpy as jnp
import jax
from jax.numpy.linalg import norm
from jax import jit, value_and_grad, random, custom_jvp
from jax.experimental import stax
from jax.experimental.stax import Dense
from jax.nn import log_sigmoid, sigmoid
from jax.nn.initializers import normal, zeros
from jax.flatten_util import ravel_pytree

from tensorflow_probability.substrates.jax.optimizer import lbfgs_minimize

from util import register_pytree_node_dataclass, tree_zeros_like
from dataclasses import dataclass
import dataclasses
from typing import Dict, Union, Any, Optional, Tuple

from functools import partial

Dataset = Tuple[jnp.ndarray, jnp.ndarray]
def _zeros_64(key, shape): return zeros(key, shape, dtype=jnp.dtype('float64'))


_init_params, _decision_function = stax.serial(
            Dense(1, W_init=_zeros_64, b_init=_zeros_64)
        )


def _init_random_params(rng: jnp.ndarray, shape: Any, sigma: float):
    """Initializer for randomized coefficients of the linear perturbation term"""
    params = []
    r1, r2 = random.split(rng)
    norm_init = normal(stddev=sigma, dtype=jnp.dtype('float64'))
    coef = norm_init(r1, (shape[-1], 1))
    intercept = norm_init(r2, (1,))
    param = coef, intercept
    params.append(param)
    return shape, params


def _init_data_weights(data: Dataset, data_weights: Optional[jnp.ndarray]) -> jnp.ndarray:
    """Helper function that returns initialized sample weights. If not specified, they are initialized to one."""
    if data_weights is None:
        return jnp.ones(data[0].shape[0], dtype=data[0].dtype)
    else:
        return data_weights


@partial(jit, static_argnums=(2,))
def _init_gram_matrix(inputs: jnp.ndarray, weights: jnp.ndarray, add_intercept: bool = True):
    """Compute the gram matrix for the inputs

    Args:
        inputs: 2d array of shape (n_samples, n_features) which contains the input features.
        weights: 1d array of shape (n_samples,) which contains non-negative weights for each sample.
        add_intercept: whether to append an extra constant feature to incorporate an intercept. Defaults to True.

    Returns:
        a 2d array of shape (n_features, n_features) if add_intercept is False, otherwise a 2d array of shape
        (n_features + 1, n_features + 1).
    """
    inputs_weighted = inputs * weights[:, None]
    gram_matrix = inputs_weighted.T @ inputs_weighted

    # Account for intercept
    if add_intercept:
        inputs_colsums = jnp.sum(inputs_weighted, axis=0)
        gram_matrix = jnp.block([[gram_matrix, inputs_colsums.reshape(-1, 1)],
                                 [inputs_colsums, jnp.sum(weights)]])

    return gram_matrix


@jit
def _objective(params: Any, random_params: Any, data: Dataset, lamb: float, pos_label: Union[str, int, float],
               data_weights: Optional[jnp.ndarray]) -> float:
    """Objective function for L2-regularized logistic regression. Includes a randomized perturbation term to
    improve privacy."""
    data_weights = _init_data_weights(data, data_weights)

    inputs = jnp.atleast_2d(data[0])
    targets = jnp.atleast_1d(data[1])
    targets = targets.ravel()  # ensure targets is a 1d vector
    scores = _decision_function(params, inputs).ravel()
    cross_entropy_terms = jnp.where(targets == pos_label, log_sigmoid(scores), log_sigmoid(-scores))
    cross_entropy_loss = -jnp.dot(data_weights, cross_entropy_terms)
    
    # Regularization penalty term
    coef, intercept = params[0]
    # Guo et al. (2020) penalize the coefficients *and* the bias
    l2_penalty = 0.5 * lamb * jnp.sum(data_weights) * (jnp.sum(coef**2) + jnp.sum(intercept**2))

    # Random linear term for indistinguishability/privacy
    b_coef, b_intercept = random_params[0]
    # Note: this term is divided by data_size in
    # https://github.com/facebookresearch/certified-removal/blob/2a3aa66e85b95d659830f944f644daaf41f63167/test_removal.py#L109
    rand_linear = (b_coef * coef).sum() + (b_intercept * intercept).sum()
    
    return cross_entropy_loss + l2_penalty + rand_linear


@partial(jit, static_argnums=3)
def _gradient(model: 'BinaryLogReg', data: Dataset, data_weights: Optional[jnp.ndarray],
              perturbation_term: bool) -> jnp.ndarray:
    """Ravelled gradient of the objective function with respect to the model parameters"""
    params_flat, unravel = ravel_pytree(model.params)
    random_params = model.random_params
    if not perturbation_term:
        random_params = tree_zeros_like(random_params)
    g = jax.grad(lambda p: _objective(unravel(p), random_params, data, model.lamb, model.pos_label, data_weights))
    return g(params_flat)


@jit
def _hessian(model: 'BinaryLogReg', data: Dataset, data_weights: Optional[jnp.ndarray]) -> jnp.ndarray:
    """Ravelled Hessian matrix of the objective function with respect to the model parameters"""
    params_flat, unravel = ravel_pytree(model.params)
    random_params = model.random_params
    h = jax.hessian(lambda p: _objective(unravel(p), random_params, data, model.lamb, model.pos_label, data_weights))
    return h(params_flat)


@custom_jvp
def _fit(init_params: Any, random_params: Any, data: Dataset, lamb: float, pos_label: Union[str, int, float],
         data_weights: jnp.ndarray, tolerance: float, max_iterations: int) -> Tuple[Any, Dict[str, Any]]:
    """Implementation of the `fit` method for BinaryLogReg"""
    # L-BFGS implementation in TFP is not compatible with PyTrees, so need to 
    # ravel parameters. Call ravelled parameters "position".
    init_position, unravel = ravel_pytree(init_params)
    
    # Wrapper for objective function which accepts position as argument.
    def obj(position):
        params = unravel(position)
        return _objective(params, random_params, data, lamb, pos_label, data_weights)
    
    obj_and_grad = value_and_grad(obj)

    def solve(f, init_position):
        return lbfgs_minimize(f, init_position, tolerance=tolerance, max_iterations=max_iterations)
    
    result = solve(obj_and_grad, init_position)
    params = unravel(result.position)

    return params, {'converged': result.converged, 
                    'num_iterations': result.num_iterations, 
                    'grad_norm': jnp.linalg.norm(result.objective_gradient, ord=2)}


# Custom implementation of jvp based on implicit function theorem
@_fit.defjvp
def _fit_jvp(primals, tangents):
    init_params, random_params, data, lamb, pos_label, data_weights, tolerance, max_iterations = primals
    _, _, data_dot, _, _, _, _, _ = tangents
    
    params, diagnostics = _fit(init_params, random_params, data, lamb, pos_label, data_weights, tolerance,
                               max_iterations)
    
    inputs, targets = data
    inputs_dot, _ = data_dot

    scores = _decision_function(params, inputs).ravel()
    preds = sigmoid(scores)
    preds_var = preds * (1 - preds)
    scores_dot = _decision_function(params, inputs_dot).ravel()
    
    first_term = jnp.einsum('i,i,ij->j', scores_dot, preds_var, inputs)
    first_term = jnp.concatenate((jnp.atleast_1d(inputs.shape[0]), first_term))
    second_term = jnp.einsum('i,ij->j', preds - targets, inputs_dot)
    second_term = jnp.concatenate((jnp.atleast_1d(0), second_term))
    jac_X_grad = first_term + second_term
    
    params_flat, unravel = ravel_pytree(params)
    hess = jax.hessian(lambda p: _objective(unravel(p), random_params, data, lamb, pos_label, data_weights))(params_flat)
    hess_inv = jnp.linalg.inv(hess)
    
    # Pre-multiplying by the inverse Jacobian of the gradient 
    # of the loss w.r.t. w (i.e. the inverse Hessian) gives 
    # the required result 
    tangent_out = unravel(hess_inv @ jac_X_grad), {k: jnp.zeros_like(jnp.asarray(v)) for k, v in diagnostics.items()}
    
    primal_out = params, diagnostics
    
    return primal_out, tangent_out


_fit_jit = jax.jit(_fit)


@partial(jit, static_argnums=5)
def _unlearn(model: 'BinaryLogReg', data: Dataset, delete: jnp.ndarray, retain: jnp.ndarray,
             prev_gram_matrix: Optional[jnp.ndarray], use_full_data_hess_approx: bool) -> \
        Tuple[jnp.ndarray, float, jnp.ndarray]:
    """Implementation of the `unlearn` method for BinaryLogReg"""
    inputs, _ = data
    
    # Gradient contribution due to deleted data
    grad_del = BinaryLogReg.gradient(model, data, delete, False)
    
    # Correct error in the quadratic approximation used by Guo et al. (2020). Namely, the Hessian should be
    # approximated on the objective prior to removal.
    if use_full_data_hess_approx:
        retain = jnp.logical_or(delete, retain)
    
    # Update the Gram matrix.
    if prev_gram_matrix is None:
        # Compute for the first time
        gram_matrix = _init_gram_matrix(inputs, retain, add_intercept=True)
    else:
        gram_matrix = prev_gram_matrix
        if not use_full_data_hess_approx:
            delta_gram_matrix = _init_gram_matrix(inputs, delete, add_intercept=True)
            gram_matrix = gram_matrix - delta_gram_matrix
    
    # Although the Hessian can be expressed as a sum over training instances, we can't update it incrementally. This
    # is because the model parameters change at each iteration.
    hess = BinaryLogReg.hessian(model, data, retain)

    hess_inv_grad_del = jnp.linalg.inv(hess) @ grad_del
    inputs_x_hess_inv_grad_del = (inputs * retain[:, None]) @ hess_inv_grad_del[0:-1] + hess_inv_grad_del[-1]
    delta_grnb = (0.25 * jnp.sqrt(norm(gram_matrix, ord=2)) * norm(hess_inv_grad_del, ord=2) *
                  norm(inputs_x_hess_inv_grad_del, ord=2))
    
    # Easier to perform a Newton update using ravelled parameters
    params_flat, unravel = ravel_pytree(model.params)
    params = unravel(params_flat + hess_inv_grad_del)
    
    # Update GRNB
    grnb = model.grnb + delta_grnb

    return params, grnb, gram_matrix


@register_pytree_node_dataclass
@dataclass(frozen=True)
class BinaryLogReg(CRModel):
    """Logistic regression model with support for certified removal

    Attributes:
        lamb: L2 regularization strength.
        params: model parameters, represented as a tuple containing a vector of coefficients and an intercept.
        random_params: random parameters, represented as a tuple containing a vector of coefficients and an
            intercept.
        epsilon: parameter for the certified removal guarantee, typically on the unit interval.
        delta: parameter for the certified removal guarantee, typically of order 1/(training set size).
        sigma: standard deviation of the random parameters.
        pos_label: label used to encode positive instances.
        neg_label: label used to encode negative instances.
        grnb: cumulative gradient residual norm bound.
        grnb_thres: threshold on `grnb` above which full retraining is triggered.
    """
    lamb: float = 1.0
    params: Optional[Tuple[jnp.ndarray, jnp.ndarray]] = None
    random_params: Optional[Tuple[jnp.ndarray, jnp.ndarray]] = None
    epsilon: float = 1.0
    delta: float = 1e-4
    sigma: float = 1.0
    pos_label: Union[str, int, float] = 1
    neg_label: Union[str, int, float] = 0
    grnb: float = 0.0
    
    @property
    def grnb_thres(self) -> float:
        return self.sigma * self.epsilon / jnp.sqrt(2 * jnp.log(1.5 / self.delta))

    def decision_function(self, inputs: jnp.ndarray) -> jnp.ndarray:
        """Predict confidence scores

        The confidence score for a sample is proportional to the signed distance of that sample to the hyperplane.

        Args:
            inputs: sample inputs to score. Must be a 2d array of shape (n_samples, n_features).

        Returns:
            array of shape (n_samples,) containing the confidence scores (for the positive class) for each sample.
            A positive score means the positive class would be predicted.
        """
        assert self.params is not None, "model is not fitted"
        return _decision_function(self.params, inputs)

    def predict(self, inputs: jnp.ndarray) -> jnp.ndarray:
        scores = self.decision_function(inputs)
        binary = (scores > 0)
        return jnp.where(binary, self.pos_label, self.neg_label)

    def predict_proba(self, inputs: jnp.ndarray) -> jnp.ndarray:
        """Probability estimates

        Args:
            inputs: samples inputs on which to estimate predicted probabilities. Must be a 2d array of shape
                (n_samples, n_features).

        Returns:
            array of shape (n_samples,) containing the predicted probabilities that each sample is in the positive
            class.
        """
        scores = self.decision_function(inputs)
        return sigmoid(scores)

    def fit(self, data: Dataset, rng: jnp.ndarray, data_weights: Optional[jnp.ndarray] = None,
            tolerance: float = 1e-4, max_iterations: int = 1000) -> Tuple['BinaryLogReg', Dict]:
        """Fit the model

        Args:
            data: training data represented as a tuple, where the first element is the inputs and the second element
                is the targets (labels).
            rng: a PRNGKey (an array with shape (2,) and dtype uint32).
            data_weights: non-negative weights associated with the samples in `data`.
            tolerance: gradient tolerance for L-BFGS. If the sup-norm of the gradient vector is below this value,
                L-BFGS is terminated.
            max_iterations: maximum number of iterations before L-BFGS is terminated.

        Returns:
            a tuple where the first element is the fitted model, and the second element contains diagnostics
            associated with L-BFGS.
        """
        data_weights = _init_data_weights(data, data_weights)

        inputs, _ = data
        
        # Update random parameters
        rng1, rng2 = random.split(rng)
        _, init_params = _init_params(rng1, (-1, inputs.shape[1]))
        _, random_params = _init_random_params(rng2, (-1, inputs.shape[1]), self.sigma)

        params, diagnostics = _fit_jit(init_params, random_params, data, self.lamb, self.pos_label,
                                       data_weights, tolerance, max_iterations)

        # Update model
        model_dict = dataclasses.asdict(self)
        model_dict['params'] = params
        model_dict['random_params'] = random_params
        model_dict['grnb'] = 0.0
        model = BinaryLogReg(**model_dict)

        return model, diagnostics

    def unlearn(self, data: Dataset, delete: jnp.ndarray, retain: jnp.ndarray, rng: jnp.ndarray,
                tolerance: float = 1e-4, max_iterations: int = 100, prev_gram_matrix: Optional[jnp.ndarray] = None,
                enforce_grnb_constraint: bool = True, use_full_data_hess_approx: bool = True,
                warn_retrain: bool = False) -> Tuple['BinaryLogReg', Optional[Dict], Optional[jnp.ndarray], bool]:
        """Unlearn training data

        Args:
            data: training data represented as a tuple, where the first element is the inputs and the second element
                is the targets (labels).
            delete: boolean array specifying samples in data to unlearn/delete.
            retain: boolean array specifying samples in data to keep post-deletion.
            rng: a PRNGKey (an array with shape (2,) and dtype uint32).
            tolerance: gradient tolerance for L-BFGS. If the sup-norm of the gradient vector is below this value,
                L-BFGS is terminated.
            max_iterations: maximum number of iterations before L-BFGS is terminated.
            prev_gram_matrix: gram matrix for the training data inputs prior to unlearning. If not provided, it is
                computed from scratch.
            enforce_grnb_constraint: boolean specifying whether to retrain from scratch if the cumulative gradient
                residual norm exceeds the bound that guarantees (epsilon, delta)-CR.
            use_full_data_hess_approx: boolean specifying whether to use the full training data (prior to deletion)
                to approximate the Hessian. If False, only the retained data (post-deletion) is used to approximate
                the Hessian as in Guo et al. (2020).
            warn_retrain: boolean specifying whether to emit a warning if retraining from scratch is triggered.
            
        Returns:
            a tuple where the first element is the unlearned model, the second element contains diagnostics associated
            with L-BFGS for retraining, the third element is the updated gram matrix, and the fourth element
            is a boolean indicating whether retraining was triggered.
        """
        assert self.params is not None, "model is not fitted"
        params, grnb, gram_matrix = _unlearn(self, data, delete, retain, prev_gram_matrix, use_full_data_hess_approx)
        
        def retrain(_):
            if warn_retrain:
                warnings.warn("GRNB of {} exceeds upper bound of {}. Retraining from "
                              "scratch".format(grnb, self.grnb_thres))
            model_new, diagnostics = BinaryLogReg.fit(self, data, rng, data_weights=retain, tolerance=tolerance,
                                                      max_iterations=max_iterations)
            return model_new, diagnostics, gram_matrix, True
    
        def done(_):
            # Update model
            model_dict = dataclasses.asdict(self)
            model_dict['params'] = params
            model_dict['grnb'] = grnb
            model_new = BinaryLogReg(**model_dict)
            diagnostics = {'converged': jnp.zeros((), dtype=bool), 'num_iterations': jnp.zeros((), dtype=jnp.int32),
                           'grad_norm': jnp.zeros((), dtype=jnp.float64)}
            return model_new, diagnostics, gram_matrix, False
        
        return jax.lax.cond(
            jnp.any(grnb > self.grnb_thres) & enforce_grnb_constraint,
            retrain,
            done,
            None
        )
    
    def set_sigma(self, value: float) -> 'BinaryLogReg':
        """Set sigma parameter
        
        Args:
            value: new value of sigma
        
        Returns:
            new model with sigma parameter updated
        """
        return dataclasses.replace(self, sigma=value, params=None, random_params=None, grnb=0.0)

    def objective(self, data: Dataset, data_weights: Optional[jnp.ndarray] = None) -> float:
        """Compute the objective on the given training data"""
        assert self.params is not None, "model is not fitted"
        return _objective(self.params, self.random_params, data, self.lamb, self.pos_label, data_weights)

    def gradient(self, data: Dataset, data_weights: Optional[jnp.ndarray] = None,
                 perturbation_term: bool = True) -> jnp.ndarray:
        """Compute the (ravelled) gradient of the objective with respect to the model parameters on the given
        training data"""
        assert self.params is not None, "model is not fitted"
        return _gradient(self, data, data_weights, perturbation_term)

    def hessian(self, data: Dataset, data_weights: Optional[jnp.ndarray] = None) -> jnp.ndarray:
        """Compute the (ravelled) Hessian of the objective with respect to the model parameters on the given training
        data"""
        assert self.params is not None, "model is not fitted"
        return _hessian(self, data, data_weights)

Initializer for randomized coefficients of the linear perturbation term Helper function that returns initialized sample weights. If not specified, they are initialized to one. Compute the gram matrix for the inputs      Args:         inputs: 2d array of shape (n_samples, n_features) which contains the input features.         weights: 1d array of shape (n_samples,) which contains non-negative weights for each sample.         add_intercept: whether to append an extra constant feature to incorporate an intercept. Defaults to True.      Returns:         a 2d array of shape (n_features, n_features) if add_intercept is False, otherwise a 2d array of shape         (n_features + 1, n_features + 1). # Account for intercept Objective function for L2-regularized logistic regression. Includes a randomized perturbation term to     improve privacy. # ensure targets is a 1d vector # Regularization penalty term # Guo et al. (2020) penalize the coefficients *and* the bias # Random linear term for indistinguishability/privacy # Note: this term is divided by data_size in # https://github.com/facebookresearch/certified-removal/blob/2a3aa66e85b95d659830f944f644daaf41f63167/test_removal.py#L109 Ravelled gradient of the objective function with respect to the model parameters Ravelled Hessian matrix of the objective function with respect to the model parameters Implementation of the `fit` method for BinaryLogReg # L-BFGS implementation in TFP is not compatible with PyTrees, so need to # ravel parameters. Call ravelled parameters "position". # Wrapper for objective function which accepts position as argument. # Custom implementation of jvp based on implicit function theorem # Pre-multiplying by the inverse Jacobian of the gradient # of the loss w.r.t. w (i.e. the inverse Hessian) gives # the required result Implementation of the `unlearn` method for BinaryLogReg # Gradient contribution due to deleted data # Correct error in the quadratic approximation used by Guo et al. (2020). Namely, the Hessian should be # approximated on the objective prior to removal. # Update the Gram matrix. # Compute for the first time # Although the Hessian can be expressed as a sum over training instances, we can't update it incrementally. This # is because the model parameters change at each iteration. # Easier to perform a Newton update using ravelled parameters # Update GRNB Logistic regression model with support for certified removal      Attributes:         lamb: L2 regularization strength.         params: model parameters, represented as a tuple containing a vector of coefficients and an intercept.         random_params: random parameters, represented as a tuple containing a vector of coefficients and an             intercept.         epsilon: parameter for the certified removal guarantee, typically on the unit interval.         delta: parameter for the certified removal guarantee, typically of order 1/(training set size).         sigma: standard deviation of the random parameters.         pos_label: label used to encode positive instances.         neg_label: label used to encode negative instances.         grnb: cumulative gradient residual norm bound.         grnb_thres: threshold on `grnb` above which full retraining is triggered. Predict confidence scores          The confidence score for a sample is proportional to the signed distance of that sample to the hyperplane.          Args:             inputs: sample inputs to score. Must be a 2d array of shape (n_samples, n_features).          Returns:             array of shape (n_samples,) containing the confidence scores (for the positive class) for each sample.             A positive score means the positive class would be predicted. Probability estimates          Args:             inputs: samples inputs on which to estimate predicted probabilities. Must be a 2d array of shape                 (n_samples, n_features).          Returns:             array of shape (n_samples,) containing the predicted probabilities that each sample is in the positive             class. Fit the model          Args:             data: training data represented as a tuple, where the first element is the inputs and the second element                 is the targets (labels).             rng: a PRNGKey (an array with shape (2,) and dtype uint32).             data_weights: non-negative weights associated with the samples in `data`.             tolerance: gradient tolerance for L-BFGS. If the sup-norm of the gradient vector is below this value,                 L-BFGS is terminated.             max_iterations: maximum number of iterations before L-BFGS is terminated.          Returns:             a tuple where the first element is the fitted model, and the second element contains diagnostics             associated with L-BFGS. # Update random parameters # Update model Unlearn training data          Args:             data: training data represented as a tuple, where the first element is the inputs and the second element                 is the targets (labels).             delete: boolean array specifying samples in data to unlearn/delete.             retain: boolean array specifying samples in data to keep post-deletion.             rng: a PRNGKey (an array with shape (2,) and dtype uint32).             tolerance: gradient tolerance for L-BFGS. If the sup-norm of the gradient vector is below this value,                 L-BFGS is terminated.             max_iterations: maximum number of iterations before L-BFGS is terminated.             prev_gram_matrix: gram matrix for the training data inputs prior to unlearning. If not provided, it is                 computed from scratch.             enforce_grnb_constraint: boolean specifying whether to retrain from scratch if the cumulative gradient                 residual norm exceeds the bound that guarantees (epsilon, delta)-CR.             use_full_data_hess_approx: boolean specifying whether to use the full training data (prior to deletion)                 to approximate the Hessian. If False, only the retained data (post-deletion) is used to approximate                 the Hessian as in Guo et al. (2020).             warn_retrain: boolean specifying whether to emit a warning if retraining from scratch is triggered.                      Returns:             a tuple where the first element is the unlearned model, the second element contains diagnostics associated             with L-BFGS for retraining, the third element is the updated gram matrix, and the fourth element             is a boolean indicating whether retraining was triggered. # Update model Set sigma parameter                  Args:             value: new value of sigma                  Returns:             new model with sigma parameter updated Compute the objective on the given training data Compute the (ravelled) gradient of the objective with respect to the model parameters on the given         training data Compute the (ravelled) Hessian of the objective with respect to the model parameters on the given training         data