class_docstr: Applies a Bijector to the diagonal of a matrix.
  
  #### Example
  
  ```python
  b = tfb.TransformDiagonal(diag_bijector=tfb.Exp())
  
  b.forward([[1., 0.],
             [0., 1.]])
  # ==> [[2.718, 0.],
         [0., 2.718]]
  ```
class_docstr_tok: ["Applies","a","Bijector","to","the","diagonal","of","a","matrix",".","#","#","#","#","Example","``","`","python","b","=","tfb.TransformDiagonal","(","diag_bijector=tfb.Exp","(",")",")","b.forward","(","[","[","1.",",","0",".","]",",","[","0.",",","1",".","]","]",")","#","==",">","[","[","2.718",",","0",".","]",",","[","0.",",","2.718","]","]","``","`"]
code: def _forward_log_det_jacobian(self, x):
      # We formulate the Jacobian with respect to the flattened matrices
      # `vec(x)` and `vec(y)`. Suppose for notational convenience that
      # the first `n` entries of `vec(x)` are the diagonal of `x`, and
      # the remaining `n**2-n` entries are the off-diagonals in
      # arbitrary order. Then the Jacobian is a block-diagonal matrix,
      # with the Jacobian of the diagonal bijector in the first block,
      # and the identity Jacobian for the remaining entries (since this
      # bijector acts as the identity on non-diagonal entries):
      #
      # J_vec(x) (vec(y)) =
      # -------------------------------
      # | J_diag(x) (diag(y))      0  | n entries
      # |                             |
      # | 0                        I  | n**2-n entries
      # -------------------------------
      #   n                     n**2-n
      #
      # Since the log-det of the second (identity) block is zero, the
      # overall log-det-jacobian is just the log-det of first block,
      # from the diagonal bijector.
      #
      # Note that for elementwise operations (exp, softplus, etc) the
      # first block of the Jacobian will itself be a diagonal matrix,
      # but our implementation does not require this to be true.
      return self.diag_bijector.forward_log_det_jacobian(
          tf.linalg.diag_part(x), event_ndims=1
      )
  
file_path: probability/tensorflow_probability/python/bijectors/transform_diagonal.py
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import bijector
label_desc: Use image node_id 6 for calling the TransformDiagonal obj's underlying member method code with example usage: obj._forward_log_det_jacobian(x) and returns: self