Custom theano Op to do numerical integration

Questions
22

Looks correct. You can set a flag self._func.trust_input=True after you first define _func. That will skip checking the type of the inputs which can cause quise some overhead when calling the function.

If you are interested we could add this as a utility directly in PyTensor (our current backend/fork of Aesara).

We would probably tweak it slightly in that case. Usually we use FunctionGraph to represent inner graphs and only compile as a function later on. There is an example of a new Loop Op that we are building you can have a look at: Implement new Loop and Scan operators by ricardoV94 ¡ Pull Request #191 ¡ pymc-devs/pytensor ¡ GitHub (specifically look at pytensor/loop/op.py::Loop)

23

Thank you for the help! I tried the advice re trust_input.

Adding self._func.trust_input=True in __init__ right after the definition of self._func makes the kernel die when running the example in a jupyter notebook. When running the code from the command line, it raises a Segmentation fault.

Any additional ideas on how to make it faster would also be great!

24

That means some of the inputs are not of the right type

25

I am struggling to see what the problem could be - a, b, and t all seem to be the right type to me in the example above. I tried with scalar b and it still raises the Segmentation error. Any thoughts / hints?

26

Can you print their types (and dtypes) in the perform method (before you pass them to _func)? That should be pretty clear. Maybe you have python literals you need to convert to numpy arrays.

27

I added:

print(inputs)
print([type(a) for a in inputs])
print([a.dtype for a in inputs])

to perform, and when running the example above the output is:

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
[array(1.), array(2.), array(2.08049756), array([5.46187651, 5.20988144, 4.94399615])]
[<class 'numpy.ndarray'>, <class 'numpy.ndarray'>, <class 'numpy.ndarray'>, <class 'numpy.ndarray'>]
[dtype('float64'), dtype('float64'), dtype('float64'), dtype('float64')]
Segmentation fault

The inputs are of the expected shapes and correspond to start, stop, a, and b.

Unless I am missing something, this looks as it should, so I guess the problem must be somewhere else (?).

28

What are the node input types? If setting trust input to false is causing a segfault it must be that the function was converting them to other types (maybe float32?)

29

Everything seems to be right:

Adding this to __init__:

print("Input types: ", [i.dtype for i in [var] + list(extra_vars)])
print("Output type: ", self._expr.type, self._expr.dtype)

prints out:

Input types:  ['float64', 'float64', 'float64']
Output type:  TensorType(float64, (?,)) float64

I also added this to make_node:

print("Make node input types: ", [
    (i.type, i.dtype)
    for i in
    [self._start, self._stop] + list(self._extra_vars_node)
])
print(
    "Make node output type: ", 
    self._expr.type()
)

which prints:

Make node input types:  [(TensorType(float64, ()), 'float64'), (TensorType(float64, ()), 'float64'), (TensorType(float64, ()), 'float64'), (TensorType(float64, (3,)), 'float64')]
Make node output type:  <TensorType(float64, (?,))>

I thought the problem could be that it doesn’t know the exact output shape, but passing it directly doesn’t seem to solve the problem either.

30

Hmm, by this point I would have to go through the self._func.__call__ and see what change is happening when trust_input is True. Anyway, that’s likely a rabbit hole not worth pursuing :slight_smile: