idata = run_analysis(y, X, Z)
File "/Users/Library/Application Support/JetBrains/PyCharmCE2023.2/scratches/scratch_LMM3.py", line 119, in run_analysis
Λ += pt.linalg.matrix_dot(X_i.T, iV_i, X_i)
HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.
```
I guess I need to rewrite the for-loop section using `scan`, but so far I haven’t been able to figure out how to do it… The point where I’m stuck is when I try to pass the list of PyTensor matrices, `basic.py` tries to convert the list to a NumPy array despite each element of the list being a different shape of matrices. Below is my attempt and the resultant error message. Any advice would be hugely appreciated!
```python
offset = 0
X_list = [None]*C
n_list = np.ndarray(shape=C)
for i in range(C):
n = sum(Z[:, i] == 1)
X_list[i] = pt.as_tensor_variable(X[offset:(offset + n), :])
n_list[i] = n
offset += n
n_list = pt.as_tensor_variable(n_list)
def oneStep(X_i, n, Λ_tm, a, b):
iV_i = 1/a * pt.eye(n) - b/(a*(a + b*n)) * pt.ones(shape=[n, n])
Λ = Λ_tm + pt.linalg.matrix_dot(X_i.T, iV_i, X_i)
return Λ
Λ_ini = pytensor.shared(np.zeros([d, d]))
Λ = pytensor.scan(fn=oneStep, outputs_info=Λ_ini,
sequences=[X_list, n_list],
non_sequences=[a, b])
```
```python
Traceback (most recent call last):
File "/Users/Library/Application Support/JetBrains/PyCharmCE2023.2/scratches/scratch_LMM3.py", line 141, in run_analysis
Λ = pytensor.scan(fn=oneStep, outputs_info=Λ_ini,
File "/Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/site-packages/pytensor/scan/basic.py", line 595, in scan
_seq_val = pt.as_tensor_variable(seq["input"])
File "/Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/site-packages/pytensor/tensor/ __init__.py", line 50, in as_tensor_variable
return _as_tensor_variable(x, name, ndim, **kwargs)
File "/Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/functools.py", line 889, in wrapper
return dispatch(args[0]. __class__ )(*args, **kw)
File "/Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/site-packages/pytensor/tensor/basic.py", line 177, in _as_tensor_Sequence
return constant(x, name=name, ndim=ndim, dtype=dtype)
File "/Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/site-packages/pytensor/tensor/basic.py", line 223, in constant
x_ = ps.convert(x, dtype=dtype)
File "/Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/site-packages/pytensor/scalar/basic.py", line 267, in convert
x_ = np.asarray(x)
ValueError: setting an array element with a sequence. The requested array has an inhomogeneous shape after 1 dimensions. The detected shape was (64,) + inhomogeneous part.
```
---
### Author:  [@yOZ](https://discourse.pymc.io/u/yOZ)
#### Post date: [July 9, 2024, 1:18am UTC](https://discourse.pymc.io/t/for-loop-does-not-work-for-customdist-distribution/14742/2 "2024-07-09T01:18:23Z")
I’ve made some updates. I could have rewritten the for-loop part using `scan` as shown below. However, now I’ve encountered a performance issue. This code involves dynamic slicing, so I cannot use the JAX/NumPyro sampler (it throws a `NotImplementedError: JAX does not support slicing arrays with a dynamic slice length.`). When I tried to test this `scan` version and the covariance matrix with `pm.MvNormal` as in the first snippet, but using `pymc.sample()`, it took 60-120 minutes, whereas the for-loop version with NumPyro only took 3 minutes.
```auto
offset = 0
n_list = np.zeros(C)
offset_list = np.zeros(C)
for i in range(C):
n = sum(Z[:, i] == 1)
n_list[i] = n
offset_list[i] = offset
offset += n
def oneStep(n, offset, Λ_tm, X, a, b):
X_i = X[offset:(offset + n), :]
iV_i = 1/a * pt.eye(n) - b/(a*(a + b*n)) * pt.ones(shape=[n, n])
Λ = Λ_tm + pt.linalg.matrix_dot(X_i.T, iV_i, X_i)
return Λ
Λ_ini = pt.as_tensor_variable(np.zeros([d, d]))
n_tv = pt.as_tensor_variable(n_list.astype('int32'))
offset_tv = pt.as_tensor_variable(offset_list.astype('int32'))
Λ, _ = pytensor.scan(fn=oneStep, outputs_info=Λ_ini,
sequences=[n_tv, offset_tv],
non_sequences=[pt.as_tensor_variable(X), a, b],
strict=True)
Λ = Λ[-1]
```
So, let me summarize the points I want to ask. The situation is that I want to use my custom distribution with `CustomDist` and feed the covariance matrix involving a for-loop to construct the custom distribution. I’d also like to use the NumPyro sampler instead of `pymc.sample()` because it’s much faster.
1. The for-loop causes an `AttributeError: 'Scratchpad' object has no attribute 'ufunc'` when combined with `CustomDist` (see the first post). How should I resolve this?
2. I tried rewriting the for-loop by `scan` (see this post). It runs without errors, but it’s very slow and doesn’t allow me to use JAX/NumPyro. Are there more efficient ways to write it?
Thank you in advance for your advice!
---
### Author:  [@ricardoV94](https://discourse.pymc.io/u/ricardoV94)
#### Post date: [July 9, 2024, 5:26am UTC](https://discourse.pymc.io/t/for-loop-does-not-work-for-customdist-distribution/14742/3 "2024-07-09T05:26:15Z")
If you are using the latest PyMC you can try calling [pymc.model.transform.optimization.freeze\_dims\_and\_data — PyMC dev documentation](https://www.pymc.io/projects/docs/en/latest/api/model/generated/pymc.model.transform.optimization.freeze_dims_and_data.html) and then sampling with numpyro.
Otherwise you can try the nutpie sampler which uses the numba backend.
Are you sure your implementations are the same? Numpyro sampler can also be very fast when diverging, which is useless.
The error you were getting was probably because your unrolled loop (without scan) had too many nodes and PyTensor fails to the function. I’ve seen that elsewhere but haven’t been able to reproduce.
---
### Author:  [@ricardoV94](https://discourse.pymc.io/u/ricardoV94)
#### Post date: [July 9, 2024, 5:28am UTC](https://discourse.pymc.io/t/for-loop-does-not-work-for-customdist-distribution/14742/4 "2024-07-09T05:28:04Z")
If you could provide a reproducible snippet for the ufunc part that would be useful for us to finally reproduce that error
---
### Author:  [@yOZ](https://discourse.pymc.io/u/yOZ)
#### Post date: [July 10, 2024, 6:08am UTC](https://discourse.pymc.io/t/for-loop-does-not-work-for-customdist-distribution/14742/5 "2024-07-10T06:08:44Z")
Thank you very much for your advice! Unfortunately, `freeze_dims_and_data()` with numpyro didn’t work somehow, as it still produced the same error (`NotImplementedError: JAX does not support slicing arrays with a dynamic slice length.`). However, nutpie sampler worked and the speed was decent. 😊
Yes, I noticed that numpyro finishes very fast when diverging, but I suppose that was not the case here, as I either see no divergence warning or just a few divergent samples (usually less than 10) even when it happens.
I have attached my snippet in case you want to take a look.
[Non-local\_LMM.py](https://discourse.pymc.io/uploads/short-url/j2EAZEBtzQjLJ7Gayu3XTKov6bP.py) (6.8 KB)
Thank you again for your help!