code: def _sparse_or_dense_matmul_onehot(sparse_or_dense_matrix, col_index):
      """Returns a (dense) column of a Tensor or SparseTensor.
  
      Args:
        sparse_or_dense_matrix: matrix-shaped, `float` `Tensor` or `SparseTensor`.
        col_index: scalar, `int` `Tensor` representing the index of the desired
          column.
  
      Returns:
        column: vector-shaped, `float` `Tensor` with the same dtype as
          `sparse_or_dense_matrix`, representing the `col_index`th column of
          `sparse_or_dense_matrix`.
      """
      if isinstance(
          sparse_or_dense_matrix,
          (tf.SparseTensor, tf1.SparseTensorValue),
      ):
          # TODO(b/111924846): Implement better (ideally in a way that allows us to
          # eliminate the `num_rows` arg, if possible).
          num_rows = _get_shape(sparse_or_dense_matrix)[-2]
          batch_shape = _get_shape(sparse_or_dense_matrix)[:-2]
          slice_start = tf.concat(
              [tf.zeros_like(batch_shape), [0, col_index]], axis=0
          )
          slice_size = tf.concat([batch_shape, [num_rows, 1]], axis=0)
          # We momentarily lose static shape information in tf.sparse_slice. However
          # we regain it in the following tf.reshape.
          sparse_slice = tf.sparse.slice(
              sparse_or_dense_matrix,
              tf.cast(slice_start, tf.int64),
              tf.cast(slice_size, tf.int64),
          )
  
          output_shape = tf.concat([batch_shape, [num_rows]], axis=0)
          return tf.reshape(
              tf.sparse.to_dense(sparse_slice), output_shape
          )
      else:
          return tf.gather(sparse_or_dense_matrix, col_index, axis=-1)
  
file_path: probability/tensorflow_probability/python/optimizer/proximal_hessian_sparse.py
imports: import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.math.generic import soft_threshold
  from tensorflow_probability.python.math.linalg import sparse_or_dense_matvecmul
label_desc: Use image node_id 7 for calling a global function with example usage: _sparse_or_dense_matmul_onehot(sparse_or_dense_matrix, col_index) and returns: tf, tf