def write_ragged_to_checkpoint(var_name, sp_mat, checkpoint_path):
  """Write scipy CSR matrix to checkpoint for loading to ragged tensor."""
  data = sp_mat.data
  indices = sp_mat.indices
  rowsplits = sp_mat.indptr
  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_rowsplits = tf.py_func(
        lambda: rowsplits.astype(np.int64), [], tf.int64, stateful=False)
    init_rowsplits.set_shape(rowsplits.shape)
    tf_data = tf.get_variable(var_name + "_data", initializer=init_data)
    tf_indices = tf.get_variable(
        var_name + "_indices", initializer=init_indices)
    tf_rowsplits = tf.get_variable(
        var_name + "_rowsplits", initializer=init_rowsplits)
    saver = tf.train.Saver([tf_data, tf_indices, tf_rowsplits])
    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.shape[0]) + " " + str(sp_mat.getnnz()))
Write scipy CSR matrix to checkpoint for loading to ragged tensor.
language/labs/drkit/search_utils.py