def make_keras_optimizer_class(cls):
    'Constructs a DP Keras optimizer class from an existing one.'

    class DPOptimizerClass(cls):
        'Differentially private subclass of given class cls.\n\n    The class tf.keras.optimizers.Optimizer has two methods to compute\n    gradients, `_compute_gradients` and `get_gradients`. The first works\n    with eager execution, while the second runs in graph mode and is used\n    by canned estimators.\n\n    Internally, DPOptimizerClass stores hyperparameters both individually\n    and encapsulated in a `GaussianSumQuery` object for these two use cases.\n    However, this should be invisible to users of this class.\n    '

        def __init__(self, l2_norm_clip, noise_multiplier, changing_clipping=False, num_microbatches=None, gradient_norm=None, *args, **kwargs):
            'Initialize the DPOptimizerClass.\n\n      Args:\n        l2_norm_clip: Clipping norm (max L2 norm of per microbatch gradients)\n        noise_multiplier: Ratio of the standard deviation to the clipping norm\n        num_microbatches: The number of microbatches into which each minibatch\n          is split.\n      '
            super(DPOptimizerClass, self).__init__(*args, **kwargs)
            self._l2_norm_clip = l2_norm_clip
            self._norm_clip = tf.Variable(l2_norm_clip)
            self._noise_multiplier = noise_multiplier
            self._num_microbatches = num_microbatches
            self._dp_sum_query = gaussian_query.GaussianSumQuery(l2_norm_clip, (l2_norm_clip * noise_multiplier))
            self._global_state = None
            self._was_dp_gradients_called = False
            self._changing_clipping = changing_clipping
            self.gradient_norm = gradient_norm

        def _compute_gradients(self, loss, var_list, grad_loss=None, tape=None):
            'DP version of superclass method.'
            self._was_dp_gradients_called = True
            if ((not callable(loss)) and (tape is None)):
                raise ValueError('`tape` is required when a `Tensor` loss is passed.')
            tape = (tape if (tape is not None) else tf.GradientTape())
            if callable(loss):
                with tape:
                    if (not callable(var_list)):
                        tape.watch(var_list)
                    if callable(loss):
                        loss = loss()
                        microbatch_losses = tf.reduce_mean(tf.reshape(loss, [self._num_microbatches, (- 1)]), axis=1)
                    if callable(var_list):
                        var_list = var_list()
            else:
                with tape:
                    microbatch_losses = tf.reduce_mean(tf.reshape(loss, [self._num_microbatches, (- 1)]), axis=1)
            var_list = tf.nest.flatten(var_list)
            with tf.keras.backend.name_scope((self._name + '/gradients')):
                jacobian = tape.jacobian(microbatch_losses, var_list)
                if self._changing_clipping:
                    if False:
                        self._norm_clip.assign_add(0.002)
                        tf.print('cur C:', self._norm_clip.value(), output_stream=sys.stdout)
                    else:
                        gr_norm = jacobian.copy()
                        for i in range(len(gr_norm)):
                            gr_norm[i] = tf.norm(gr_norm[i])
                        gr_mean = tf.math.reduce_mean(gr_norm)
                        self._norm_clip.assign(gr_mean)
                        C = self._norm_clip.value()
                        tf.print('cur C:', C, output_stream=sys.stdout)

                def clip_gradients(g):
                    return tf.clip_by_global_norm(g, self._norm_clip.value())[0]
                clipped_gradients = tf.map_fn(clip_gradients, jacobian)

                def reduce_noise_normalize_batch(g):
                    summed_gradient = tf.reduce_sum(g, axis=0)
                    noise_stddev = (self._l2_norm_clip * self._noise_multiplier)
                    noise = tf.random.normal(tf.shape(input=summed_gradient), stddev=noise_stddev)
                    noised_gradient = tf.add(summed_gradient, noise)
                    return tf.truediv(noised_gradient, self._num_microbatches)
                final_gradients = tf.nest.map_structure(reduce_noise_normalize_batch, clipped_gradients)
            return list(zip(final_gradients, var_list))

        def get_gradients(self, loss, params):
            if self._changing_clipping:
                self._l2_norm_clip *= 0.99
                tf.print('cur C:', self._l2_norm_clip, output_stream=sys.stdout)
            'DP version of superclass method.'
            self._was_dp_gradients_called = True
            if (self._global_state is None):
                self._global_state = self._dp_sum_query.initial_global_state()
            microbatch_losses = tf.reshape(loss, [self._num_microbatches, (- 1)])
            sample_params = self._dp_sum_query.derive_sample_params(self._global_state)

            def process_microbatch(i, sample_state):
                'Process one microbatch (record) with privacy helper.'
                mean_loss = tf.reduce_mean(input_tensor=tf.gather(microbatch_losses, [i]))
                grads = tf.gradients(mean_loss, params)
                sample_state = self._dp_sum_query.accumulate_record(sample_params, sample_state, grads)
                return sample_state
            sample_state = self._dp_sum_query.initial_sample_state(params)
            for idx in range(self._num_microbatches):
                sample_state = process_microbatch(idx, sample_state)
            (grad_sums, self._global_state) = self._dp_sum_query.get_noised_result(sample_state, self._global_state)

            def normalize(v):
                try:
                    return tf.truediv(v, tf.cast(self._num_microbatches, tf.float32))
                except TypeError:
                    return None
            final_grads = tf.nest.map_structure(normalize, grad_sums)
            return final_grads

        def apply_gradients(self, grads_and_vars, global_step=None, name=None):
            assert self._was_dp_gradients_called, 'Neither _compute_gradients() or get_gradients() on the differentially private optimizer was called. This means the training is not differentially private. It may be the case that you need to upgrade to TF 2.4 or higher to use this particular optimizer.'
            return super(DPOptimizerClass, self).apply_gradients(grads_and_vars, global_step, name)
    return DPOptimizerClass
Constructs a DP Keras optimizer class from an existing one.
def make_keras_optimizer_class(cls):
    

    class DPOptimizerClass(cls):
        'Differentially private subclass of given class cls.\n\n    The class tf.keras.optimizers.Optimizer has two methods to compute\n    gradients, `_compute_gradients` and `get_gradients`. The first works\n    with eager execution, while the second runs in graph mode and is used\n    by canned estimators.\n\n    Internally, DPOptimizerClass stores hyperparameters both individually\n    and encapsulated in a `GaussianSumQuery` object for these two use cases.\n    However, this should be invisible to users of this class.\n    '

        def __init__(self, l2_norm_clip, noise_multiplier, changing_clipping=False, num_microbatches=None, gradient_norm=None, *args, **kwargs):
            'Initialize the DPOptimizerClass.\n\n      Args:\n        l2_norm_clip: Clipping norm (max L2 norm of per microbatch gradients)\n        noise_multiplier: Ratio of the standard deviation to the clipping norm\n        num_microbatches: The number of microbatches into which each minibatch\n          is split.\n      '
            super(DPOptimizerClass, self).__init__(*args, **kwargs)
            self._l2_norm_clip = l2_norm_clip
            self._norm_clip = tf.Variable(l2_norm_clip)
            self._noise_multiplier = noise_multiplier
            self._num_microbatches = num_microbatches
            self._dp_sum_query = gaussian_query.GaussianSumQuery(l2_norm_clip, (l2_norm_clip * noise_multiplier))
            self._global_state = None
            self._was_dp_gradients_called = False
            self._changing_clipping = changing_clipping
            self.gradient_norm = gradient_norm

        def _compute_gradients(self, loss, var_list, grad_loss=None, tape=None):
            'DP version of superclass method.'
            self._was_dp_gradients_called = True
            if ((not callable(loss)) and (tape is None)):
                raise ValueError('`tape` is required when a `Tensor` loss is passed.')
            tape = (tape if (tape is not None) else tf.GradientTape())
            if callable(loss):
                with tape:
                    if (not callable(var_list)):
                        tape.watch(var_list)
                    if callable(loss):
                        loss = loss()
                        microbatch_losses = tf.reduce_mean(tf.reshape(loss, [self._num_microbatches, (- 1)]), axis=1)
                    if callable(var_list):
                        var_list = var_list()
            else:
                with tape:
                    microbatch_losses = tf.reduce_mean(tf.reshape(loss, [self._num_microbatches, (- 1)]), axis=1)
            var_list = tf.nest.flatten(var_list)
            with tf.keras.backend.name_scope((self._name + '/gradients')):
                jacobian = tape.jacobian(microbatch_losses, var_list)
                if self._changing_clipping:
                    if False:
                        self._norm_clip.assign_add(0.002)
                        tf.print('cur C:', self._norm_clip.value(), output_stream=sys.stdout)
                    else:
                        gr_norm = jacobian.copy()
                        for i in range(len(gr_norm)):
                            gr_norm[i] = tf.norm(gr_norm[i])
                        gr_mean = tf.math.reduce_mean(gr_norm)
                        self._norm_clip.assign(gr_mean)
                        C = self._norm_clip.value()
                        tf.print('cur C:', C, output_stream=sys.stdout)

                def clip_gradients(g):
                    return tf.clip_by_global_norm(g, self._norm_clip.value())[0]
                clipped_gradients = tf.map_fn(clip_gradients, jacobian)

                def reduce_noise_normalize_batch(g):
                    summed_gradient = tf.reduce_sum(g, axis=0)
                    noise_stddev = (self._l2_norm_clip * self._noise_multiplier)
                    noise = tf.random.normal(tf.shape(input=summed_gradient), stddev=noise_stddev)
                    noised_gradient = tf.add(summed_gradient, noise)
                    return tf.truediv(noised_gradient, self._num_microbatches)
                final_gradients = tf.nest.map_structure(reduce_noise_normalize_batch, clipped_gradients)
            return list(zip(final_gradients, var_list))

        def get_gradients(self, loss, params):
            if self._changing_clipping:
                self._l2_norm_clip *= 0.99
                tf.print('cur C:', self._l2_norm_clip, output_stream=sys.stdout)
            'DP version of superclass method.'
            self._was_dp_gradients_called = True
            if (self._global_state is None):
                self._global_state = self._dp_sum_query.initial_global_state()
            microbatch_losses = tf.reshape(loss, [self._num_microbatches, (- 1)])
            sample_params = self._dp_sum_query.derive_sample_params(self._global_state)

            def process_microbatch(i, sample_state):
                'Process one microbatch (record) with privacy helper.'
                mean_loss = tf.reduce_mean(input_tensor=tf.gather(microbatch_losses, [i]))
                grads = tf.gradients(mean_loss, params)
                sample_state = self._dp_sum_query.accumulate_record(sample_params, sample_state, grads)
                return sample_state
            sample_state = self._dp_sum_query.initial_sample_state(params)
            for idx in range(self._num_microbatches):
                sample_state = process_microbatch(idx, sample_state)
            (grad_sums, self._global_state) = self._dp_sum_query.get_noised_result(sample_state, self._global_state)

            def normalize(v):
                try:
                    return tf.truediv(v, tf.cast(self._num_microbatches, tf.float32))
                except TypeError:
                    return None
            final_grads = tf.nest.map_structure(normalize, grad_sums)
            return final_grads

        def apply_gradients(self, grads_and_vars, global_step=None, name=None):
            assert self._was_dp_gradients_called, 'Neither _compute_gradients() or get_gradients() on the differentially private optimizer was called. This means the training is not differentially private. It may be the case that you need to upgrade to TF 2.4 or higher to use this particular optimizer.'
            return super(DPOptimizerClass, self).apply_gradients(grads_and_vars, global_step, name)
    return DPOptimizerClass