Adding a custom sparse op (Sparse Determinant) - c++

Question

I am working on trying to get some sparse matrix operations working in Tensorflow. The first one I am tackling is a sparse determinant, via a sparse Cholesky decomposition. Eigen has a sparse Cholesky, so my thought is to wrap that.
I have been making some progress, but am now a little bit stuck. I know that SparseTensors in Tensorflow are made up of three parts: indices, values, and shape. Copying similar ops, I went for the following REGISTER_OP declaration:
REGISTER_OP("SparseLogDet")
.Input("a_indices: int64")
.Input("a_values: float32")
.Input("a_shape: int64")
.Output("determinant: float32")
.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
shape_inference::ShapeHandle h;
c->set_output(0, h);
return Status::OK();
});
This compiles fine, but when I run it using some example code:
import tensorflow as tf
log_det_op = tf.load_op_library('./sparse_log_det_op.so')
with tf.Session(''):
t = tf.SparseTensor(indices=[[0, 0], [1, 2]], values=[1, 2],
dense_shape=[3, 4])
print(log_det_op.sparse_log_det(t).eval().shape)
print(log_det_op.sparse_log_det(t).eval())
It complains, saying:
TypeError: sparse_log_det() missing 2 required positional arguments: 'a_values' and 'a_shape'
This makes sense to me, since it's expecting the other arguments. However, I would really just like to pass the sparse tensor, not break it up into components! Does anyone know how this is handled for other sparse operations?
Thanks!

Answer 1

If you want to pass in the sparse tensor and then determine indices, values and shape from this, this should be possible. Just modify your OP to take a single Tensor input, and produce a single float output. Then extract the desired information form the Eigen::Tensor by looping through its elements as seen below:
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/shape_inference.h"
#include "tensorflow/core/framework/op_kernel.h"
#include <Eigen/Dense>
using namespace tensorflow;
REGISTER_OP("SparseDeterminant")
.Input("sparse_tensor: float")
.Output("sparse_determinant: float");
class SparseDeterminantOp : public OpKernel {
public:
explicit SparseDeterminantOp(OpKernelConstruction *context) : OpKernel(context) {}
void Compute(OpKernelContext *context) override {
// get the input tesnorflow tensor
const Tensor& sparse_tensor = context->input(0);
// get shape of input
const TensorShape& sparse_shape = sparse_tensor.shape();
// get Eigen Tensor for input tensor
auto eigen_sparse = sparse_tensor.matrix<float>();
//extract the data you want from the sparse tensor input
auto a_shape = sparse_tensor.shape();
// loop over all elements of the input tensor and add to values and indices
for (int i=0; i<a_shape.dim_size(0); ++i){
for (int j=0; j<a_shape.dim_size(1); ++j){
if(eigen_sparse(i,j) != 0){
/// ***Here add non zero elements to list/tensor of values and their indicies***
std::cout<<eigen_sparse(i,j)<<" at"<<" "<<i<<" "<<j<<" "<<"not zero."<<std::endl;
}
}
}
// create output tensor
Tensor *output_tensor = NULL;
TensorShape output_shape;
OP_REQUIRES_OK(context, context->allocate_output(0, output_shape, &output_tensor));
auto output = output_tensor->scalar<float>();
output(0) = 1.; //**asign return value***;
}
};
REGISTER_KERNEL_BUILDER(Name("SparseDeterminant").Device(DEVICE_CPU), SparseDeterminantOp);
sadly, when you pass t into your op it becomes a Tensorflow::Tensor and loses the values and indices methods associated with tf.sparsetensor, so you can't get them easily.
Once compiled this code can be run with:
//run.py
import tensorflow as tf
import numpy as np
my_module = tf.load_op_library('./out.so')
# create sparse matrix
a = np.zeros((10,10))
for i in range(len(a)):
a[i,i] = i
print(a)
a_t = tf.convert_to_tensor(a, dtype= float)
with tf.Session() as sess:
sess.run(my_module.sparse_determinant(a_t))