code: def _one_nonzero_per_row(matrix):
      """For each row in matrix, randomly zero all but one of the nonzero values."""
      # TODO(b/172564129): can it be done without a loop?
      out = np.zeros_like(matrix)
      for i in range(matrix.shape[0]):
          nonzero_indices = np.flatnonzero(matrix[i])
          keep = np.random.choice(nonzero_indices)
          out[i, keep] = matrix[i, keep]
      return out
  
file_path: model_search/model_search/search/linear_model.py
imports: from model_search import block_builder
  from model_search.architecture import architecture_utils
  from model_search.proto import phoenix_spec_pb2
  from model_search.search import common
  from model_search.search import search_algorithm
  import numpy
  import sklearn
  import sklearn.linear_model
label_desc: Use image node_id 1 for calling a global function with example usage: _one_nonzero_per_row(matrix) and returns: out