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 vectorizes the computation of a clipped gradient defined over a single sample.
  
  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","vectorizes","the","computation","of","a","clipped","gradient","defined","over","a","single","sample",".","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_simple(
      self, f: Callable, gv: GradValues, vc: VarCollection
  ) -> Callable:
      """Creates a function that computes gradients clipped per-sample.
      This algorithm vectorizes the computation of a clipped gradient defined over a single sample.
  
      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 clipped_grad_single_example(*args):
          grads, values = gv(*args)
          total_grad_norm = jn.linalg.norm(
              jn.array([jn.linalg.norm(g) for g in grads])
          )
          idivisor = 1 / jn.maximum(
              total_grad_norm / self.l2_norm_clip, 1.0
          )
          return [g * idivisor for g in grads], values
  
      clipped_grad_vectorized = Vectorize(
          clipped_grad_single_example, batch_axis=self.batch_axis
      )
  
      def clipped_grad(*args):
          g, v = clipped_grad_vectorized(*args)
          g, v = jax.tree_map(
              functools.partial(jn.mean, axis=0), (g, v)
          )
          return g, v
  
      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 2 for calling the PrivateGradValues obj's underlying member method code with example usage: obj._make_clipped_grad_fn_simple(f, gv, vc) and returns: clipped_grad, unknown,  values, g,  v