Error using sparse array

1

Hello

I’m running into an error running the following code:

import numpy as np
import pymc as pm
import pytensor as pt

for module in pm, pt, np:
    print(module.__name__, module.__version__)

np.random.seed(42)

arr = np.random.choice([0, 1], size=(250, 100), p=[0.9, 0.1])

arr_sparse = pt.sparse.csr_from_dense(pt.tensor.as_tensor(arr))

with pm.Model():
    beta = pm.Normal("beta", 0, 1, shape=arr.shape[1])
    mu = pt.sparse.dot(arr_sparse, beta)
    sigma = pm.Exponential("sigma", 1)
    pm.Normal("lik", mu, sigma, observed=np.random.randn(arr.shape[0]))
    pm.sample()

The error:

pymc 5.16.2
pytensor 2.25.2
numpy 1.26.4
Traceback (most recent call last):
  File "~/micromamba/envs/pymc_mm_env/lib/python3.12/site-packages/pytensor/tensor/math.py", line 1861, in dot
    res = l.__dot__(r)
          ^^^^^^^^^^^^
  File "~/micromamba/envs/pymc_mm_env/lib/python3.12/site-packages/pytensor/sparse/basic.py", line 390, in __dot__
    return structured_dot(left, right)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "~/micromamba/envs/pymc_mm_env/lib/python3.12/site-packages/pytensor/sparse/basic.py", line 3562, in structured_dot
    return _structured_dot(x, y)
           ^^^^^^^^^^^^^^^^^^^^^
  File "~/micromamba/envs/pymc_mm_env/lib/python3.12/site-packages/pytensor/graph/op.py", line 293, in __call__
    node = self.make_node(*inputs, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "~/micromamba/envs/pymc_mm_env/lib/python3.12/site-packages/pytensor/sparse/basic.py", line 3442, in make_node
    raise NotImplementedError("non-matrix b")
NotImplementedError: non-matrix b

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "~/djpiri/content/AgScreens/H5/barda-gates/2024-02-16-db-regression/sparse.py", line 20, in <module>
    pm.sample()
  File "~/micromamba/envs/pymc_mm_env/lib/python3.12/site-packages/pymc/sampling/mcmc.py", line 716, in sample
    step = assign_step_methods(model, step, methods=pm.STEP_METHODS, step_kwargs=kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "~/micromamba/envs/pymc_mm_env/lib/python3.12/site-packages/pymc/sampling/mcmc.py", line 223, in assign_step_methods
    tg.grad(model_logp, var)  # type: ignore
    ^^^^^^^^^^^^^^^^^^^^^^^^
  File "~/micromamba/envs/pymc_mm_env/lib/python3.12/site-packages/pytensor/gradient.py", line 607, in grad
    _rval: Sequence[Variable] = _populate_grad_dict(
                                ^^^^^^^^^^^^^^^^^^^^
  File "~/micromamba/envs/pymc_mm_env/lib/python3.12/site-packages/pytensor/gradient.py", line 1400, in _populate_grad_dict
    rval = [access_grad_cache(elem) for elem in wrt]
            ^^^^^^^^^^^^^^^^^^^^^^^
  File "~/micromamba/envs/pymc_mm_env/lib/python3.12/site-packages/pytensor/gradient.py", line 1355, in access_grad_cache
    term = access_term_cache(node)[idx]
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "~/micromamba/envs/pymc_mm_env/lib/python3.12/site-packages/pytensor/gradient.py", line 1185, in access_term_cache
    input_grads = node.op.L_op(inputs, node.outputs, new_output_grads)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "~/micromamba/envs/pymc_mm_env/lib/python3.12/site-packages/pytensor/graph/op.py", line 398, in L_op
    return self.grad(inputs, output_grads)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "~/micromamba/envs/pymc_mm_env/lib/python3.12/site-packages/pytensor/sparse/basic.py", line 4028, in grad
    rval.append(pt_dot(x.T, gz))
                ^^^^^^^^^^^^^^^
  File "~/micromamba/envs/pymc_mm_env/lib/python3.12/site-packages/pytensor/tensor/math.py", line 1865, in dot
    res = r.__rdot__(l)
          ^^^^^^^^^^^^^
  File "~/micromamba/envs/pymc_mm_env/lib/python3.12/site-packages/pytensor/tensor/variable.py", line 654, in __rdot__
    return pt.math.dense_dot(left, right)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "~/micromamba/envs/pymc_mm_env/lib/python3.12/site-packages/pytensor/tensor/math.py", line 1914, in dense_dot
    raise TypeError("The dense dot product is only supported for dense types")
TypeError: The dense dot product is only supported for dense types

It looks like a dense implementation of the dot product is getting called at some point. Presumably this isn’t intended?

It’s also complaining about beta not being a matrix. Calling mu.eval() before pm.sample() works fine suggesting beta is fine as it is.

Any help would be great.

Many thanks
David

1 Like
2

Hey! I had this problem in the past, I think you need:

mu = pytensor.sparse.dot(arr_sparse, beta[:, None]).flatten()

The reason is that the dot product expects two matrices, but your beta is a vector. The result will be a matrix, but you need a vector, so you flatten.

1 Like
3

Thanks! You just sped up my sampling 30x!

1 Like
4

It’s a bit confusing that pymc.math.dot and pytensor.sparse.dot don’t behave the same (you can pass a vector to pymc.math.dot just fine).

Is this worth opening an issue to suggest?

5

Yeah, you can open an issue.

There are a lot of things we could do to the sparse module (maintenance, improvements, and additions). But nobody had a chance to do that yet. I was able to write some improvements for a project but I never found a moment to contribute that back to PyMC… I take this as a friendly reminder :slight_smile:

6

OK. I just started a discussion: pytensor.sparse.dot could accept vectors like pymc.math.dot · pymc-devs/pytensor · Discussion #965 · GitHub

1 Like