code: def pad_shape_with_ones(x, ndims, start=-1):
      """Maybe add `ndims` ones to `x.shape` starting at `start`.
  
      If `ndims` is zero, this is a no-op; otherwise, we will create and return a
      new `Tensor` whose shape is that of `x` with `ndims` ones concatenated on the
      right side. If the shape of `x` is known statically, the shape of the return
      value will be as well.
  
      Args:
        x: The `Tensor` we'll return a reshaping of.
        ndims: Python `integer` number of ones to pad onto `x.shape`.
        start: Python `integer` specifying where to start padding with ones. Must
          be a negative integer. For instance, a value of `-1` means to pad at the
          end of the shape. Default value: `-1`.
      Returns:
        If `ndims` is zero, `x`; otherwise, a `Tensor` whose shape is that of `x`
        with `ndims` ones concatenated on the right side. If possible, returns a
        `Tensor` whose shape is known statically.
      Raises:
        ValueError: if `ndims` is not a Python `integer` greater than or equal to
        zero.
      """
      if not (isinstance(ndims, int) and ndims >= 0):
          raise ValueError(
              "`ndims` must be a Python `integer` greater than zero. Got: {}".format(
                  ndims
              )
          )
      if not (isinstance(start, int) and start <= -1):
          raise ValueError(
              "`start` must be a Python `integer` less than zero. Got: {}".format(
                  start
              )
          )
      if ndims == 0:
          return x
      x = tf.convert_to_tensor(value=x)
      original_shape = x.shape
      rank = ps.rank(x)
      first_shape = ps.shape(x)[: rank + start + 1]
      second_shape = ps.shape(x)[rank + start + 1 :]
      new_shape = ps.pad(
          first_shape, paddings=[[0, ndims]], constant_values=1
      )
      new_shape = ps.concat([new_shape, second_shape], axis=0)
      x = tf.reshape(x, new_shape)
      if start == -1:
          tensorshape_util.set_shape(
              x,
              tensorshape_util.concatenate(original_shape, [1] * ndims),
          )
      elif tensorshape_util.rank(original_shape) is not None:
          original_ndims = tensorshape_util.rank(original_shape)
          new_shape = tensorshape_util.concatenate(
              original_shape[: original_ndims + start + 1],
              tensorshape_util.concatenate(
                  [1] * ndims,
                  original_shape[original_ndims + start + 1 :],
              ),
          )
          tensorshape_util.set_shape(x, new_shape)
      return x
  
file_path: probability/tensorflow_probability/python/math/psd_kernels/internal/util.py
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
label_desc: Use image node_id 1 for calling a global function with example usage: pad_shape_with_ones(x, ndims, start) and returns: x, x