class_docstr: Base class for various modules which compute derivatives.
class_docstr_tok: ["Base","class","for","various","modules","which","compute","derivatives","."]
docstr: Constructs an instance to compute the derivatives of f w.r.t. variables.
  Args:
      derivative_fn: JAX transformation which computes derivative.
      f: the function for which to compute derivatives.
      variables: the variables for which to compute derivatives.
      input_argnums: input indexes, if any, on which to compute derivatives.
      return_all_f_outputs: if True also return original outputs of the fuction along with derivatives.
docstr_tok: ["Constructs","an","instance","to","compute","the","derivatives","of","f","w.r.t",".","variables",".","Args",":","derivative_fn",":","JAX","transformation","which","computes","derivative",".","f",":","the","function","for","which","to","compute","derivatives",".","variables",":","the","variables","for","which","to","compute","derivatives",".","input_argnums",":","input","indexes",",","if","any",",","on","which","to","compute","derivatives",".","return_all_f_outputs",":","if","True","also","return","original","outputs","of","the","fuction","along","with","derivatives","."]
code: def __init__(
      self,
      derivative_fn: Callable,
      f: Union[Module, Callable],
      variables: Optional[VarCollection],
      input_argnums: Optional[Tuple[int, ...]] = None,
      return_all_f_outputs: bool = False,
  ):
      """Constructs an instance to compute the derivatives of f w.r.t. variables.
      Args:
          derivative_fn: JAX transformation which computes derivative.
          f: the function for which to compute derivatives.
          variables: the variables for which to compute derivatives.
          input_argnums: input indexes, if any, on which to compute derivatives.
          return_all_f_outputs: if True also return original outputs of the fuction along with derivatives.
      """
      self.f = f
      self.vc = variables = VarCollection(variables or ())
      if not isinstance(f, Module):
          f = Function(f, self.vc)
  
      assert (
          isinstance(input_argnums, tuple) or input_argnums is None
      ), f"Must pass a tuple of indices to input_argnums; received {input_argnums}."
      self.input_argnums = input_argnums or tuple()
      self.return_all_f_outputs = return_all_f_outputs
  
      def f_func(
          inputs_and_train_tensors: List[JaxArray],
          state_tensors: List[JaxArray],
          list_args: List,
          kwargs: Dict,
      ):
          inputs = inputs_and_train_tensors[: len(self.input_argnums)]
          train_tensors = inputs_and_train_tensors[
              len(self.input_argnums) :
          ]
          original_vc = self.vc.tensors()
          try:
              self.vc.subset(TrainVar).assign(train_tensors)
              self.vc.subset(BaseState).assign(state_tensors)
              for i, arg in zip(self.input_argnums, inputs):
                  list_args[i] = arg
              outputs = f(*list_args, **kwargs)
              if not isinstance(outputs, (list, tuple)):
                  outputs = [outputs]
              if self.return_all_f_outputs:
                  return outputs[0], (outputs, variables.tensors())
              else:
                  return outputs[0], variables.tensors()
          finally:
              self.vc.assign(original_vc)
  
      self._call = derivative_fn(f_func)
  
imports: import inspect
  from typing import List, Optional, Callable, Tuple, Dict, Union
  import jax
  from objax.module import Function, Module
  from objax.typing import JaxArray
  from objax.util import repr_function, class_name
  from objax.variable import BaseState, TrainVar, VarCollection
label_desc: Use image node_id 1 to create a new _DerivativeBase object from inherited base classes: Module with example: obj = _DerivativeBase(derivative_fn, f, variables, input_argnums, return_all_f_outputs)