def write_sparse_to_checkpoint(var_name, sp_mat, checkpoint_path):
  """Write scipy sparse CSR matrix to checkpoint."""
  sp_mat = sp_mat.tocoo()
  # Sort the indices lexicographically.
  sort_i = np.lexsort((sp_mat.col, sp_mat.row))
  indices = np.asmatrix([sp_mat.row[sort_i], sp_mat.col[sort_i]]).transpose()
  data = sp_mat.data[sort_i]
  with tf.Graph().as_default():
    init_data = tf.py_func(
        lambda: data.astype(np.float32), [], tf.float32, stateful=False)
    init_data.set_shape(data.shape)
    init_indices = tf.py_func(
        lambda: indices.astype(np.int64), [], tf.int64, stateful=False)
    init_indices.set_shape(indices.shape)
    init_shape = tf.py_func(
        lambda: np.array(sp_mat.shape, dtype=np.int64), [],
        tf.int64,
        stateful=False)
    init_shape.set_shape([len(sp_mat.shape)])
    tf_data = tf.get_variable(var_name + "_data", initializer=init_data)
    tf_indices = tf.get_variable(
        var_name + "_indices", initializer=init_indices)
    tf_shape = tf.get_variable(var_name + "_shape", initializer=init_shape)
    saver = tf.train.Saver([tf_data, tf_indices, tf_shape])
    with tf.Session() as session:
      session.run(tf.global_variables_initializer())
      saver.save(session, checkpoint_path)
  with tf.gfile.Open(checkpoint_path + ".info", "w") as f:
    f.write(str(sp_mat.getnnz()))
Write scipy sparse CSR matrix to checkpoint.
language/labs/drkit/search_utils.py