def diag_jacobian(xs,
                  ys=None,
                  sample_shape=None,
                  fn=None,
                  parallel_iterations=10,
                  name=None):
  """
  """
  with tf.compat.v1.name_scope(name, 'jacobians_diag', [xs, ys]):
    if sample_shape is None:
      sample_shape = [1]
    # Output Jacobian diagonal
    jacobians_diag_res = []
    # Convert input `xs` to a list
    xs = list(xs) if _is_list_like(xs) else [xs]
    xs = [tf.convert_to_tensor(value=x) for x in xs]
    if not tf.executing_eagerly():
      if ys is None:
        if fn is None:
          raise ValueError('Both `ys` and `fn` can not be `None`')
        else:
          ys = fn(*xs)
      # Convert ys to a list
      ys = list(ys) if _is_list_like(ys) else [ys]
      if len(xs) != len(ys):
        raise ValueError('`xs` and `ys` should have the same length')
      for y, x in zip(ys, xs):
        # Broadcast `y` to the shape of `x`.
        y_ = y + tf.zeros_like(x)
        # Change `event_shape` to one-dimension
        y_ = tf.reshape(y, tf.concat([sample_shape, [-1]], -1))

        # Declare an iterator and tensor array loop variables for the gradients.
        n = tf.size(input=x) / tf.cast(
            tf.reduce_prod(input_tensor=sample_shape), dtype=tf.int32)
        n = tf.cast(n, dtype=tf.int32)
        loop_vars = [
            0,
            tf.TensorArray(x.dtype, n)
        ]

        def loop_body(j):
          """Loop function to compute gradients of the each direction."""
          # Gradient along direction `j`.
          res = tf.gradients(ys=y_[..., j], xs=x)[0]  # pylint: disable=cell-var-from-loop
          if res is None:
            # Return zero, if the gradient is `None`.
            res = tf.zeros(tf.concat([sample_shape, [1]], -1),
                           dtype=x.dtype)  # pylint: disable=cell-var-from-loop
          else:
            # Reshape `event_shape` to 1D
            res = tf.reshape(res, tf.concat([sample_shape, [-1]], -1))
            # Add artificial dimension for the case of zero shape input tensor
            res = tf.expand_dims(res, 0)
            res = res[..., j]
          return res  # pylint: disable=cell-var-from-loop

        # Iterate over all elements of the gradient and compute second order
        # derivatives.
        _, jacobian_diag_res = tf.while_loop(
            cond=lambda j, _: j < n,  # pylint: disable=cell-var-from-loop
            body=lambda j, result: (j + 1, result.write(j, loop_body(j))),
            loop_vars=loop_vars,
            parallel_iterations=parallel_iterations)

        shape_x = tf.shape(input=x)
        # Stack gradients together and move flattened `event_shape` to the
        # zero position
        reshaped_jacobian_diag = tf.transpose(a=jacobian_diag_res.stack())
        # Reshape to the original tensor
        reshaped_jacobian_diag = tf.reshape(reshaped_jacobian_diag, shape_x)
        jacobians_diag_res.append(reshaped_jacobian_diag)

    else:
      if fn is None:
        raise ValueError('`fn` can not be `None` when eager execution is '
                         'enabled')
      if ys is None:
        ys = fn(*xs)

      def fn_slice(i, j):
        """Broadcast y[i], flatten event shape of y[i], return y[i][..., j]."""
        def fn_broadcast(*state):
          res = fn(*state)
          res = list(res) if _is_list_like(res) else [res]
          if len(res) != len(state):
            res *= len(state)
          res = [tf.reshape(r + tf.zeros_like(s),
                            tf.concat([sample_shape, [-1]], -1))
                 for r, s in zip(res, state)]
          return res
        # Expand dimensions before returning in order to support 0D input `xs`
        return lambda *state: tf.expand_dims(fn_broadcast(*state)[i], 0)[..., j]

      def make_loop_body(i, x):
        """Loop function to compute gradients of the each direction."""
        def _fn(j, result):
          res = value_and_gradient(fn_slice(i, j), xs)[1][i]
          if res is None:
            res = tf.zeros(tf.concat([sample_shape, [1]], -1), dtype=x.dtype)
          else:
            res = tf.reshape(res, tf.concat([sample_shape, [-1]], -1))
            res = res[..., j]
          return j + 1, result.write(j, res)
        return _fn

      for i, x in enumerate(xs):
        # Declare an iterator and tensor array loop variables for the gradients.
        n = tf.size(input=x) / tf.cast(
            tf.reduce_prod(input_tensor=sample_shape), dtype=tf.int32)
        n = tf.cast(n, dtype=tf.int32)
        loop_vars = [
            0,
            tf.TensorArray(x.dtype, n)
        ]

        # Iterate over all elements of the gradient and compute second order
        # derivatives.
        _, jacobian_diag_res = tf.while_loop(
            cond=lambda j, _: j < n,
            body=make_loop_body(i, x),
            loop_vars=loop_vars,
            parallel_iterations=parallel_iterations)

        shape_x = tf.shape(input=x)
        # Stack gradients together and move flattened `event_shape` to the
        # zero position
        reshaped_jacobian_diag = tf.transpose(a=jacobian_diag_res.stack())
        # Reshape to the original tensor
        reshaped_jacobian_diag = tf.reshape(reshaped_jacobian_diag, shape_x)
        jacobians_diag_res.append(reshaped_jacobian_diag)

  return ys, jacobians_diag_res