RandomType pickling issues

[ricardoV94]

This problem was seen in pymc-devs/pymc#5090

import numpy as np
import aesara
# import cloudpickle as pickle  # either library fails
import pickle

# same problem with aesara.shared(np.random.RandomState())
var1 = aesara.shared(np.random.default_rng())
var2 = aesara.shared(np.random.default_rng())
print(var2)  # RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FC441E6EAC0>)

res = var1.type.filter_variable(var2)
print(res)  # RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FC441E6EAC0>)

var2 = pickle.loads(pickle.dumps(var2))
print(var2)  # RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FC441E853C0>)

res = var1.type.filter_variable(var2)  # ! raises TypeError

TypeError: Cannot convert Type RandomGeneratorType (of Variable <RandomGeneratorType>) into Type RandomGeneratorType. You can try to manually convert <RandomGeneratorType> into a RandomGeneratorType.

[brandonwillard]

The RandomGeneratorType instances are not equal, because the pickle/un-pickle process needs to construct a new one when it builds the last instance of var2 (i.e. it doesn't have a reference to the original instance).

These Type subclasses should probably be singletons (i.e. we need to implement custom __reduce__ methods that use the singleton instances). A cheaper solution is to implement type-only checking in custom __eq__ methods on those subclasses.