code: def matrix_diag_transform(matrix, transform=None, name=None):
      """Transform diagonal of [batch-]matrix, leave rest of matrix unchanged.
  
      Create a trainable covariance defined by a Cholesky factor:
  
      ```python
      # Transform network layer into 2 x 2 array.
      matrix_values = tf.contrib.layers.fully_connected(activations, 4)
      matrix = tf.reshape(matrix_values, (batch_size, 2, 2))
  
      # Make the diagonal positive. If the upper triangle was zero, this would be a
      # valid Cholesky factor.
      chol = matrix_diag_transform(matrix, transform=tf.nn.softplus)
  
      # LinearOperatorLowerTriangular ignores the upper triangle.
      operator = LinearOperatorLowerTriangular(chol)
      ```
  
      Example of heteroskedastic 2-D linear regression.
  
      ```python
      tfd = tfp.distributions
  
      # Get a trainable Cholesky factor.
      matrix_values = tf.contrib.layers.fully_connected(activations, 4)
      matrix = tf.reshape(matrix_values, (batch_size, 2, 2))
      chol = matrix_diag_transform(matrix, transform=tf.nn.softplus)
  
      # Get a trainable mean.
      mu = tf.contrib.layers.fully_connected(activations, 2)
  
      # This is a fully trainable multivariate normal!
      dist = tfd.MultivariateNormalTriL(mu, chol)
  
      # Standard log loss. Minimizing this will "train" mu and chol, and then dist
      # will be a distribution predicting labels as multivariate Gaussians.
      loss = -1 * tf.reduce_mean(dist.log_prob(labels))
      ```
  
      Args:
        matrix:  Rank `R` `Tensor`, `R >= 2`, where the last two dimensions are
          equal.
        transform:  Element-wise function mapping `Tensors` to `Tensors`. To be
          applied to the diagonal of `matrix`. If `None`, `matrix` is returned
          unchanged. Defaults to `None`.
        name:  A name to give created ops. Defaults to "matrix_diag_transform".
  
      Returns:
        A `Tensor` with same shape and `dtype` as `matrix`.
      """
      with ops.name_scope(name, "matrix_diag_transform", [matrix]):
          matrix = ops.convert_to_tensor(matrix, name="matrix")
          if transform is None:
              return matrix
          # Replace the diag with transformed diag.
          diag = array_ops.matrix_diag_part(matrix)
          transformed_diag = transform(diag)
          transformed_mat = array_ops.matrix_set_diag(
              matrix, transformed_diag
          )
  
      return transformed_mat
  
imports: import functools
  import hashlib
  import numpy
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import tensor_shape
  from tensorflow.python.framework import tensor_util
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import array_ops_stack
  from tensorflow.python.ops import check_ops
  from tensorflow.python.ops import cond
  from tensorflow.python.ops import control_flow_ops
  from tensorflow.python.ops import linalg_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import nn
  from tensorflow.python.util import tf_inspect
label_desc: Use image node_id 16 for calling a global function with example usage: matrix_diag_transform(matrix, transform, name) and returns: transformed_mat, matrix