class_docstr: Implements a continuous normalizing flow X->Y defined via an ODE.
  
  
  This bijector implements a continuous dynamics transformation
  parameterized by a differential equation, where initial and terminal
  conditions correspond to domain (X) and image (Y) i.e.
  
  ```None
  d/dt[state(t)]=state_time_derivative_fn(t, state(t))
  state(initial_time) = X
  state(final_time) = Y
  ```
  
  For this transformation the value of `log_det_jacobian` follows another
  differential equation, reducing it to computation of the trace of the jacbian
  along the trajectory
  
  ```None
  state_time_derivative = state_time_derivative_fn(t, state(t))
  d/dt[log_det_jac(t)] = Tr(jacobian(state_time_derivative, state(t)))
  ```
  
  FFJORD constructor takes two functions `ode_solve_fn` and
  `trace_augmentation_fn` arguments that customize integration of the
  differential equation and trace estimation.
  
  Differential equation integration is performed by a call to `ode_solve_fn`.
  Custom `ode_solve_fn` must accept the following arguments:
  * ode_fn(time, state, **condition_kwargs): Differential equation to be solved.
    Custom `ode_solve_fn`s may optionally support conditional inputs by
    accepting a `constants` dict arg and computing gradients wrt the provided
    values in **condition_kwargs.
  * initial_time: Scalar float or floating Tensor representing the initial time.
  * initial_state: Floating Tensor representing the initial state.
  * solution_times: 1D floating Tensor of solution times.
  
  And return a Tensor of shape [solution_times.shape, initial_state.shape]
  representing state values evaluated at `solution_times`. In addition
  `ode_solve_fn` must support nested structures. For more details see the
  interface of `tfp.math.ode.Solver.solve()`.
  
  Trace estimation is computed simultaneously with `state_time_derivative`
  using `augmented_state_time_derivative_fn` that is generated by
  `trace_augmentation_fn`. `trace_augmentation_fn` takes
  `state_time_derivative_fn`, `state.shape` and `state.dtype` arguments and
  returns a `augmented_state_time_derivative_fn` callable that computes both
  `state_time_derivative` and unreduced `trace_estimation`.
  
  #### Custom `ode_solve_fn` and `trace_augmentation_fn` examples:
  ```python
  # custom_solver_fn: `callable(f, t_initial, t_solutions, y_initial, ...)`
  # custom_solver_kwargs: Additional arguments to pass to custom_solver_fn.
  def ode_solve_fn(ode_fn, initial_time, initial_state, solution_times):
    results = custom_solver_fn(ode_fn, initial_time, solution_times,
                               initial_state, **custom_solver_kwargs)
    return results
  
  ffjord = tfb.FFJORD(state_time_derivative_fn, ode_solve_fn=ode_solve_fn)
  ```
  
  ```python
  # state_time_derivative_fn: `callable(time, state)`
  # trace_jac_fn: `callable(time, state)` unreduced jacobian trace function
  
  def trace_augmentation_fn(ode_fn, state_shape, state_dtype):
    def augmented_ode_fn(time, state):
      return ode_fn(time, state), trace_jac_fn(time, state)
    return augmented_ode_fn
  
  ffjord = tfb.FFJORD(state_time_derivative_fn,
                      trace_augmentation_fn=trace_augmentation_fn)
  ```
  
  For more details on FFJORD and continous normalizing flows see [1], [2].
  
  #### Usage example:
  ```python
  tfd = tfp.distributions
  tfb = tfp.bijectors
  # state_time_derivative_fn: `Callable(time, state)` -> state_time_derivative
  # e.g. Neural network with inputs and outputs of the same shapes and dtypes.
  
  bijector = tfb.FFJORD(state_time_derivative_fn=state_time_derivative_fn)
  y = bijector.forward(x)  # forward mapping
  x = bijector.inverse(y)  # inverse mapping
  base = tfd.Normal(tf.zeros_like(x), tf.ones_like(x))  # Base distribution
  transformed_distribution = tfd.TransformedDistribution(base, bijector)
  ```
  
  #### References
  
  [1]:  Chen, T. Q., Rubanova, Y., Bettencourt, J., & Duvenaud, D. K. (2018).
        Neural ordinary differential equations. In Advances in neural
        information processing systems (pp. 6571-6583)
  
  [2]:  Grathwohl, W., Chen, R. T., Betterncourt, J., Sutskever, I.,
        & Duvenaud, D. (2018). Ffjord: Free-form continuous dynamics for
        scalable reversible generative models. arXiv preprint
        arXiv:1810.01367.
        http://arxiv.org.abs/1810.01367
class_docstr_tok: ["Implements","a","continuous","normalizing","flow","X-",">","Y","defined","via","an","ODE",".","This","bijector","implements","a","continuous","dynamics","transformation","parameterized","by","a","differential","equation",",","where","initial","and","terminal","conditions","correspond","to","domain","(","X",")","and","image","(","Y",")","i.e",".","``","`","None","d\/dt","[","state","(","t",")","]","=state_time_derivative_fn","(","t",",","state","(","t",")",")","state","(","initial_time",")","=","X","state","(","final_time",")","=","Y","``","`","For","this","transformation","the","value","of","`","log_det_jacobian","`","follows","another","differential","equation",",","reducing","it","to","computation","of","the","trace","of","the","jacbian","along","the","trajectory","``","`","None","state_time_derivative","=","state_time_derivative_fn","(","t",",","state","(","t",")",")","d\/dt","[","log_det_jac","(","t",")","]","=","Tr","(","jacobian","(","state_time_derivative",",","state","(","t",")",")",")","``","`","FFJORD","constructor","takes","two","functions","`","ode_solve_fn","`","and","`","trace_augmentation_fn","`","arguments","that","customize","integration","of","the","differential","equation","and","trace","estimation",".","Differential","equation","integration","is","performed","by","a","call","to","`","ode_solve_fn","`",".","Custom","`","ode_solve_fn","`","must","accept","the","following","arguments",":","*","ode_fn","(","time",",","state",",","*","*","condition_kwargs",")",":","Differential","equation","to","be","solved",".","Custom","`","ode_solve_fn","`","s","may","optionally","support","conditional","inputs","by","accepting","a","`","constants","`","dict","arg","and","computing","gradients","wrt","the","provided","values","in","*","*","condition_kwargs",".","*","initial_time",":","Scalar","float","or","floating","Tensor","representing","the","initial","time",".","*","initial_state",":","Floating","Tensor","representing","the","initial","state",".","*","solution_times",":","1D","floating","Tensor","of","solution","times",".","And","return","a","Tensor","of","shape","[","solution_times.shape",",","initial_state.shape","]","representing","state","values","evaluated","at","`","solution_times","`",".","In","addition","`","ode_solve_fn","`","must","support","nested","structures",".","For","more","details","see","the","interface","of","`","tfp.math.ode.Solver.solve","(",")","`",".","Trace","estimation","is","computed","simultaneously","with","`","state_time_derivative","`","using","`","augmented_state_time_derivative_fn","`","that","is","generated","by","`","trace_augmentation_fn","`",".","`","trace_augmentation_fn","`","takes","`","state_time_derivative_fn","`",",","`","state.shape","`","and","`","state.dtype","`","arguments","and","returns","a","`","augmented_state_time_derivative_fn","`","callable","that","computes","both","`","state_time_derivative","`","and","unreduced","`","trace_estimation","`",".","#","#","#","#","Custom","`","ode_solve_fn","`","and","`","trace_augmentation_fn","`","examples",":","``","`","python","#","custom_solver_fn",":","`","callable","(","f",",","t_initial",",","t_solutions",",","y_initial",",","...",")","`","#","custom_solver_kwargs",":","Additional","arguments","to","pass","to","custom_solver_fn",".","def","ode_solve_fn","(","ode_fn",",","initial_time",",","initial_state",",","solution_times",")",":","results","=","custom_solver_fn","(","ode_fn",",","initial_time",",","solution_times",",","initial_state",",","*","*","custom_solver_kwargs",")","return","results","ffjord","=","tfb.FFJORD","(","state_time_derivative_fn",",","ode_solve_fn=ode_solve_fn",")","``","`","``","`","python","#","state_time_derivative_fn",":","`","callable","(","time",",","state",")","`","#","trace_jac_fn",":","`","callable","(","time",",","state",")","`","unreduced","jacobian","trace","function","def","trace_augmentation_fn","(","ode_fn",",","state_shape",",","state_dtype",")",":","def","augmented_ode_fn","(","time",",","state",")",":","return","ode_fn","(","time",",","state",")",",","trace_jac_fn","(","time",",","state",")","return","augmented_ode_fn","ffjord","=","tfb.FFJORD","(","state_time_derivative_fn",",","trace_augmentation_fn=trace_augmentation_fn",")","``","`","For","more","details","on","FFJORD","and","continous","normalizing","flows","see","[","1","]",",","[","2","]",".","#","#","#","#","Usage","example",":","``","`","python","tfd","=","tfp.distributions","tfb","=","tfp.bijectors","#","state_time_derivative_fn",":","`","Callable","(","time",",","state",")","`","-",">","state_time_derivative","#","e.g",".","Neural","network","with","inputs","and","outputs","of","the","same","shapes","and","dtypes",".","bijector","=","tfb.FFJORD","(","state_time_derivative_fn=state_time_derivative_fn",")","y","=","bijector.forward","(","x",")","#","forward","mapping","x","=","bijector.inverse","(","y",")","#","inverse","mapping","base","=","tfd.Normal","(","tf.zeros_like","(","x",")",",","tf.ones_like","(","x",")",")","#","Base","distribution","transformed_distribution","=","tfd.TransformedDistribution","(","base",",","bijector",")","``","`","#","#","#","#","References","[","1","]",":","Chen",",","T.","Q.",",","Rubanova",",","Y.",",","Bettencourt",",","J.",",","&","Duvenaud",",","D.","K.","(","2018",")",".","Neural","ordinary","differential","equations",".","In","Advances","in","neural","information","processing","systems","(","pp",".","6571-6583",")","[","2","]",":","Grathwohl",",","W.",",","Chen",",","R.","T.",",","Betterncourt",",","J.",",","Sutskever",",","I.",",","&","Duvenaud",",","D.","(","2018",")",".","Ffjord",":","Free-form","continuous","dynamics","for","scalable","reversible","generative","models",".","arXiv","preprint","arXiv:1810.01367.","http",":","\/\/arxiv.org.abs\/1810.01367"]
docstr: Computes inverse and inverse_log_det_jacobian transformations.
docstr_tok: ["Computes","inverse","and","inverse_log_det_jacobian","transformations","."]
code: def _augmented_inverse(self, y, **condition_kwargs):
      """Computes inverse and inverse_log_det_jacobian transformations."""
      augmented_inv_ode_fn = self._trace_augmentation_fn(
          self._inv_state_time_derivative_fn, ps.shape(y), y.dtype
      )
      augmented_y = (y, tf.zeros_like(y))
      if condition_kwargs:
          x, ildj = self._solve_ode(
              augmented_inv_ode_fn,
              augmented_y,
              constants=condition_kwargs,
          )
      else:
          x, ildj = self._solve_ode(augmented_inv_ode_fn, augmented_y)
      return x, {"ildj": ildj, "fldj": -ildj}
  
file_path: probability/tensorflow_probability/python/bijectors/ffjord.py
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import bijector
  from tensorflow_probability.python.internal import cache_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.math.diag_jacobian import diag_jacobian
  from tensorflow_probability.python.math.ode import dormand_prince
label_desc: Use image node_id 5 for calling the FFJORD obj's underlying member method code with example usage: obj._augmented_inverse(y) and returns: x,  dict