TypeError: Real.grad illegally returned an integer-valued variable - when pm.sample

1

HI All,
I’m trying to estimate how much a mic array has been shifted using experimental measurements. Modelling the measured mic-to-mic distances for a few of the mics in the whole array - I try to check how much each mic is shifted from to match the measurements.

The problem I have is that the pm.draw works fine and generates dimensionally consistent outputs. The moment I try to pm.sample however - I get the following Error:

TypeError: Real.grad illegally  returned an integer-valued variable. (Input index 0, dtype complex128)

This is the simplest reproducible example below:

import pymc as pm
import pytensor as pt



def calc_distance_to(focalmic, mic_array):
    '''
    '''
    diff = mic_array - focalmic 
    norm_dist = pm.math.abs(pt.tensor.linalg.norm(diff, axis=1))
    return norm_dist

full_array = np.random.normal(0,10,16*3).reshape(16,3)
micnums_inds = [0,8,15] # subset only these mics each time. 
full_mic_rows = np.random.normal(0,0.1,3*16).reshape(3,16)

with pm.Model() as dev_model:
    ideal_micarray = pm.Data('ideal', full_array)
    focal_mics = pm.Data('focal_mics', micnums_inds)
    obs_micdists = pm.Data('obs_micdists', full_mic_rows)
    mean_shifts = pm.Normal('mean_shifts', 0, 0.1, shape=(full_array.shape[0],
                                                          full_array.shape[1]))
    
    sd = pm.HalfNormal('sd', 0.1)
    shift = pm.Normal('shift', mu=mean_shifts, sigma=sd)
    
    shifted_mics = pm.Deterministic('shifted_mics', ideal_micarray + shift)
    focal_micxyz = pm.Deterministic('focal_micxyz', shifted_mics[focal_mics])
    mic2mic_dists = pm.Deterministic('mic2mic_dists', pt.scan(fn=calc_distance_to,
                                                              sequences=focal_micxyz,
                            non_sequences=shifted_mics, return_updates=False))
    
    pred_sd = pm.HalfNormal('pred_sd', 0.3)
    observed_mismatch = pm.Normal('observed_mismatch', mu=mic2mic_dists,
                                  sigma=pred_sd,
                                  observed=obs_micdists)
    idata = pm.sample()

Any help or pointers would be greatly appreciated.

Using the following environment:
Python 3.11.0
Pymc 6.0.1
Pytensor 3.0.4
OS: Windows 10

2

I can now confirm the problem is the pt.scan part of the code where I use a subset of what may be a pt.dmatrix to calculate a distance matrix effectively.
When I replace the pt.scan part with a shape-conform but nonsense continuous pm.Normal - everything suddenly works.

This is the example below:

"""
Created on Sun Jul 12 13:12:43 2026

@author: theja
"""
import numpy as np 
import pymc as pm 
import scipy.spatial as spl
import pytensor as pt


full_array = np.random.normal(0,10,16*3).reshape(16,3)
shifts = np.random.normal(0,0.01, 16*3).reshape(16,3)

micnums_inds = np.array([0,8,15]) # subset only these mics each time. 
full_mic_rows = spl.distance_matrix(full_array[micnums_inds], full_array)

#%%

def calc_distance_to(focalmic, mic_array):
    '''
    '''
    diff = mic_array - focalmic
    norm_dist = pm.math.abs(pt.tensor.linalg.norm(diff, axis=1))
    return norm_dist

#%%
if __name__ == '__main__':
    with pm.Model() as dev_model:
        ideal_micarray = pm.Data('ideal', full_array)
        focal_mics = pm.Data('focal_mics', micnums_inds)
        obs_micdists = pm.Data('obs_micdists', full_mic_rows)
        mean_shifts = pm.Normal('mean_shifts', -0.01, 0.1, shape=(full_array.shape[0],
                                                              full_array.shape[1]))
        
        sd = pm.HalfNormal('sd', 0.1)
        shift = pm.Normal('shift', mu=mean_shifts, sigma=sd)
        
        shifted_mics = pm.Deterministic('shifted_mics', ideal_micarray + shift)
        focal_micxyz = shifted_mics[focal_mics,:]

        mic2mic_dists = pm.Deterministic('mic2mic_dists', pt.scan(fn=calc_distance_to,
                                                                  sequences=focal_micxyz,
                                non_sequences=shifted_mics, return_updates=False))
        # un-commenting the bttom line & commenting the above line makes everything work 
        # suddenly.
        # mic2mic_dists = pm.Normal('mic2mic_dists',1,4, shape=(3,16))
        pred_sd = pm.HalfNormal('pred_sd', 0.3)
        observed_mismatch = pm.Normal('observed_mismatch', mu=mic2mic_dists,
                                      sigma=pred_sd,
                                      observed=full_mic_rows)
        idata = pm.sample(nuts_sampler='pymc')

        ```
3

Problem solved - thanks to a related question + answer by @jessegrabowski .

The issue turned out to be the distance matrix calculation.

The function def calc_distance_to calculates the euclidean distance between a focal mic in a mic array and all other mics - leading to 0 distance for the focal mic to itself. pt.math.linalg.norm goes haywire while seeing this 0 term - and throws the error. Adding a tiny offset to the distances does not help for some reason.

Things of course worked when I replaced the focal mic indexing with a shape conform pm.Normal because the random distribution almost never produces exactly 0 values.

The solution now is somehow to add a tiny offset (e.g. 1e-9) to the distances, and using pm.math functions rather than the pt.linalg.norm. This is the version of the distance calculation that works.

def calc_distance_to(focalmic, mic_array):
    '''
    '''
    diff = mic_array - focalmic
    diff += 1e-9
    norm_dist = pm.math.sqrt(pm.math.sum(diff**2, axis=1))
    return norm_dist
1 Like