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,)