class_docstr: Computes differentially private gradients as required by DP-SGD.
  This module can be used in place of GradVals, and automatically makes the optimizer differentially private.
class_docstr_tok: ["Computes","differentially","private","gradients","as","required","by","DP-SGD",".","This","module","can","be","used","in","place","of","GradVals",",","and","automatically","makes","the","optimizer","differentially","private","."]
docstr: Creates a function that computes gradients clipped per-sample.
  This algorithm first accumulates the norm of the gradient per sample, and then
  performs an additional weighted backward pass that is equivalent to per-sample clipping.
  The accumulation allows jax to free up memory whenever the automatic differentiation engine
  is done processing a layer. This is more memory efficient and can be faster than the "simple" approach
  defined above, in particular at large batch-sizes for deep models.
  
  Reference: https://arxiv.org/abs/2009.03106.
  
  Args:
    f: the function for which to compute gradients.
    gv: the GradValues object that computes non-clipped gradients.
    vc: the variables for which to compute gradients.
  
  Returns:
    clipped_grad: the function computing the average of gradients clipped per-sample.
docstr_tok: ["Creates","a","function","that","computes","gradients","clipped","per-sample",".","This","algorithm","first","accumulates","the","norm","of","the","gradient","per","sample",",","and","then","performs","an","additional","weighted","backward","pass","that","is","equivalent","to","per-sample","clipping",".","The","accumulation","allows","jax","to","free","up","memory","whenever","the","automatic","differentiation","engine","is","done","processing","a","layer",".","This","is","more","memory","efficient","and","can","be","faster","than","the","``","simple","''","approach","defined","above",",","in","particular","at","large","batch-sizes","for","deep","models",".","Reference",":","https",":","\/\/arxiv.org\/abs\/2009.03106",".","Args",":","f",":","the","function","for","which","to","compute","gradients",".","gv",":","the","GradValues","object","that","computes","non-clipped","gradients",".","vc",":","the","variables","for","which","to","compute","gradients",".","Returns",":","clipped_grad",":","the","function","computing","the","average","of","gradients","clipped","per-sample","."]
code: def _make_clipped_grad_fn_norm_accumulation(
      self, f: Callable, gv: GradValues, vc: VarCollection
  ) -> Callable:
      """Creates a function that computes gradients clipped per-sample.
      This algorithm first accumulates the norm of the gradient per sample, and then
      performs an additional weighted backward pass that is equivalent to per-sample clipping.
      The accumulation allows jax to free up memory whenever the automatic differentiation engine
      is done processing a layer. This is more memory efficient and can be faster than the "simple" approach
      defined above, in particular at large batch-sizes for deep models.
  
      Reference: https://arxiv.org/abs/2009.03106.
  
      Args:
        f: the function for which to compute gradients.
        gv: the GradValues object that computes non-clipped gradients.
        vc: the variables for which to compute gradients.
  
      Returns:
        clipped_grad: the function computing the average of gradients clipped per-sample.
      """
  
      @Function.with_vars(gv.vars())
      def grad_norm_fn(*args):
          grads, values = gv(*args)
          total_grad_norm = jn.linalg.norm(
              jn.array([jn.linalg.norm(g) for g in grads])
          )
          return total_grad_norm, values
  
      grad_norm_fn_vectorized = Vectorize(
          grad_norm_fn, batch_axis=self.batch_axis
      )
      loss_fn_vectorized = Vectorize(
          Function.with_vars(gv.vars())(f), batch_axis=self.batch_axis
      )
  
      def weighted_loss_fn(weights, *args):
          returned = loss_fn_vectorized(*args)
          # the following assumes that the loss is the first element output by loss_fn_vectorized
          if isinstance(returned, jn.ndarray):
              loss_vector = returned
          else:
              loss_vector = returned[0]
          return jn.mean(loss_vector * weights)
  
      weighted_grad_fn = GradValues(weighted_loss_fn, vc)
  
      def clipped_grad(*args):
          grad_norms, values = grad_norm_fn_vectorized(*args)
          idivisor = 1 / jn.maximum(grad_norms / self.l2_norm_clip, 1.0)
          clipped_grads, _ = weighted_grad_fn(idivisor, *args)
          return clipped_grads, jax.tree_map(
              functools.partial(jn.mean, axis=0), values
          )
  
      return clipped_grad
  
file_path: objax/objax/privacy/dpsgd/gradient.py
imports: import functools
  from typing import Optional, Callable, Tuple
  import jax
  import jax.numpy
  from objax import random
  from objax.gradient import GradValues
  from objax.module import Function, Module, Vectorize
  from objax.typing import JaxArray
  from objax.util import repr_function, class_name
  from objax.variable import VarCollection
label_desc: Use image node_id 3 for calling the PrivateGradValues obj's underlying member method code with example usage: obj._make_clipped_grad_fn_norm_accumulation(f, gv, vc) and returns: clipped_grad, total_grad_norm,  values, jn, clipped_grads,  jax