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: Constructs a PrivateGradValues instance.
  
  Args:
      f: the function for which to compute gradients.
      vc: the variables for which to compute gradients.
      noise_multiplier: scale of standard deviation for added noise in DP-SGD.
      l2_norm_clip: value of clipping norm for DP-SGD.
      microbatch: the size of each microbatch.
      batch_axis: the axis to use as batch during vectorization. Should be a tuple of 0s.
      keygen: a Generator for random numbers. Defaults to objax.random.DEFAULT_GENERATOR.
      use_norm_accumulation: whether to use norm accumulation as a first step to compute the clipped
        gradients (more details further). This is more memory efficient and often faster when the model
        is very deep and uses a large batch-size.
docstr_tok: ["Constructs","a","PrivateGradValues","instance",".","Args",":","f",":","the","function","for","which","to","compute","gradients",".","vc",":","the","variables","for","which","to","compute","gradients",".","noise_multiplier",":","scale","of","standard","deviation","for","added","noise","in","DP-SGD",".","l2_norm_clip",":","value","of","clipping","norm","for","DP-SGD",".","microbatch",":","the","size","of","each","microbatch",".","batch_axis",":","the","axis","to","use","as","batch","during","vectorization",".","Should","be","a","tuple","of","0s",".","keygen",":","a","Generator","for","random","numbers",".","Defaults","to","objax.random.DEFAULT_GENERATOR",".","use_norm_accumulation",":","whether","to","use","norm","accumulation","as","a","first","step","to","compute","the","clipped","gradients","(","more","details","further",")",".","This","is","more","memory","efficient","and","often","faster","when","the","model","is","very","deep","and","uses","a","large","batch-size","."]
code: def __init__(
      self,
      f: Callable,
      vc: VarCollection,
      noise_multiplier: float,
      l2_norm_clip: float,
      microbatch: int,
      batch_axis: Tuple[Optional[int], ...] = (0,),
      keygen: random.Generator = random.DEFAULT_GENERATOR,
      use_norm_accumulation: bool = False,
  ):
      """Constructs a PrivateGradValues instance.
  
      Args:
          f: the function for which to compute gradients.
          vc: the variables for which to compute gradients.
          noise_multiplier: scale of standard deviation for added noise in DP-SGD.
          l2_norm_clip: value of clipping norm for DP-SGD.
          microbatch: the size of each microbatch.
          batch_axis: the axis to use as batch during vectorization. Should be a tuple of 0s.
          keygen: a Generator for random numbers. Defaults to objax.random.DEFAULT_GENERATOR.
          use_norm_accumulation: whether to use norm accumulation as a first step to compute the clipped
            gradients (more details further). This is more memory efficient and often faster when the model
            is very deep and uses a large batch-size.
      """
      super().__init__()
  
      if not all(v == 0 for v in batch_axis):
          raise ValueError(
              "batch_axis needs to be an all zero tuple for PrivateGradValues."
          )
  
      self.__wrapped__ = gv = GradValues(f, vc)
  
      self.batch_axis = batch_axis
      self.microbatch = microbatch
      self.l2_norm_clip = l2_norm_clip
      self.noise_multiplier = noise_multiplier
      self.keygen = keygen
  
      if use_norm_accumulation:
          self.clipped_grad = (
              self._make_clipped_grad_fn_norm_accumulation(f, gv, vc)
          )
      else:
          self.clipped_grad = self._make_clipped_grad_fn_simple(
              f, gv, vc
          )
  
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 1 to create a new PrivateGradValues object from inherited base classes: Module with example: obj = PrivateGradValues(f, vc, noise_multiplier, l2_norm_clip, microbatch, batch_axis, keygen, use_norm_accumulation)