import tensorflow as tf

from tensorflow.python.framework import ops

from tensorflow.python.framework import sparse_tensor
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gen_sparse_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import sparse_ops

class BatchedSpMM:
	def __init__(self):
		self.b_module = tf.load_op_library('./batched.so')

	def call(self, sp_matrices, dense_matrices, adjoint_a=False, adjoint_b=False):
		sp_indices = [sp_m.indices for sp_m in sp_matrices]
		sp_values = [sp_m.values for sp_m in sp_matrices]
		sp_shape = [sp_m.dense_shape for sp_m in sp_matrices]
		return self.b_module.bspmm(sp_ids = sp_indices, sp_values = sp_values, sp_shape = sp_shape, rhs = dense_matrices, adjoint_a = adjoint_a, adjoint_b = adjoint_b)

class BatchedSpMDT:
	def __init__(self):
