def _reduce_sum_sparse(spenv, *argspecs, axes):
X, = argspecs
data, indices, out_shape = sparse.bcoo_reduce_sum(
X.data(spenv), X.indices(spenv), shape=X.shape, axes=axes)
if out_shape == ():
out_argspec = ArgSpec(out_shape, spenv.push(data.sum()), None)
else:
out_argspec = ArgSpec(out_shape, spenv.push(data), spenv.push(indices))
return (out_argspec,)
def test_bcoo_reduce_sum(self, shape, dtype, n_batch, n_dense, axes):
rng = rand_sparse(self.rng())
M = rng(shape, dtype)
data, indices = sparse.bcoo_fromdense(M, n_batch=n_batch, n_dense=n_dense)
data_out, indices_out, shape_out = sparse.bcoo_reduce_sum(data, indices, shape=shape, axes=axes)
result_dense = M.sum(axes)
result_sparse = sparse.bcoo_todense(data_out, indices_out, shape=shape_out)
tol = {np.float32: 1E-6, np.float64: 1E-14}
self.assertAllClose(result_dense, result_sparse, atol=tol, rtol=tol)
def _reduce_sum_sparse(spenv, *spvalues, axes):
X, = spvalues
X_promoted = spvalues_to_arrays(spenv, X)
mat = sparse.bcoo_reduce_sum(X_promoted, axes=axes)
out_shape = mat.shape
if out_shape == ():
out_spvalue = spenv.dense(mat.data.sum())
else:
out_spvalue = spenv.sparse(out_shape, mat.data, mat.indices)
return (out_spvalue,)