Model metaclass for more natural model definition (without string var names on RHS)

[smartass101]

I’ve always wondered whether there might be a way to specify the model without the string name on the right-hand side like we do now

with pm.Model() as linear_model:
    weights = pm.Normal('weights', mu=0, sigma=1)   # name 'weights' specified on RHS

Recently I’ve been playing with param (mostly in the Panel context, but it’s independent) and I realized that they know the label/name of a parameter without specifying it as a string in a class constructor call because they are using a smart metaclass which sets the names according to the attribute names.

In theory this should make it possible to specify the model like so

class LinearModel(pm.Model):    # model definition as a class becomes mandatory though for this API
    weights = pm.Normal(mu=0, sigma=1)   # the name weights is encoded as the class-level attr name

if there is the following metaclass implementation (rough sketch) and names in variable constructors become optional

class ModelMetaclass(...):

    def __init__(mcs, name, bases, dict_):
        ....
        variables = ( (n, v) for (n,v) in dict_.items()
                             if isinstance(v, (Distribution, Deterministic))]   # TODO some better baseclass?
        for var_name, var in variables:
            var._set_name(var_name)   # probably has to do more things

In theory this might even be backwards compatible, it could be a next-gen API if people are interested. Personally, I encode models as a class quite often as it is convenient for stacking models, etc., so for me creating classes for models does not seems such a hassle.

One might argue that this is just syntax sugar and perhaps not that necessary. But I think that for quite a few people syntax sugar is an important decision factor (even though they may not admit it even to themselves).

I’m putting this idea here to gather feedback and see if there is any interest.

[junpenglao]

We have explored something similar, @twiecki and @_eigenfoo have a bit more context on this.

[_eigenfoo]

We’ve considered doing something similar to this in PyMC4, but not PyMC3. Even in PyMC4, the discussion was mainly around removing the yield keyword from the model specification API by parsing the Python AST, as I outlined in [1]. We thought that if we were doing this, it would make sense to also add an “autonaming” feature. However, the AST proposal was subsequently dropped [2], and we never really picked up any discussion on autonaming for PyMC4…

That doesn’t address the proposal here, though: it seems sensible enough to me, but the main difficulty is how we would get the dict_ dictionary in the @smartass101’s code snippet. I don’t think that object currently exists anywhere in the PyMC3 codebase, so we would have to construct it ourselves… by parsing the AST

Which is not to say that I’m opposed to an autonaming feature! Just want to point out that it might be nontrivial to implement.

[1] https://eigenfoo.xyz/manipulating-python-asts/
[2] https://twitter.com/avibryant/status/1150827954319982592

[smartass101]

I looked a bit deeper into the metaclass docs and I discovered that since Python 3.6 implementing this is actually trivial using the __set__name__ descriptor method. If you wanted to support older Python versions, then it would be quite easy to backport this functionality into a simple helper metaclass (a quick search on PyPi didn’t reveal such a package, but perhaps I missed it).

So on Py>=3.6 all you need to do is the following in the dsitribution base class

class Distribution:
    """Base class for distributions"""
    def __init__(self, name=None, **params):
        self.name = name
        self.params = params
        
    def __set_name__(self, owner, name):
        """Descriptor method since Py3.6
        
        Could activate other functionality depending on a known name,
        e.g. allocating nodes in the computational graph"""
        if self.name is None:  
            self.name = name
        # otherwise keep the specified one, perhaps to give nicer labels

    def __repr__(self):
        """Show a nice representation of the distribution
        
        if bound to a variable, prints is as 
        var = Distrib(params)
        """
        params = ', '.join(f'{n}={v}' for (n,v) in self.params.items())
        base = f'{type(self).__name__}({params})'
        try:
            return f'{self.name} = ' + base
        except AttributeError:
            return base

Let’s create some simple mock-up distributions for visualization purposes.

class Normal(Distribution):
    
    def __init__(self, name=None, mu=0, sd=1):
        super().__init__(name, mu=mu, sd=sd)
        
class Uniform(Distribution):
    
    def __init__(self, name=None, lower=-1, upper=1):
        super().__init__(name, lower=lower, upper=upper)

So the name argument becomes optional. Now this is what you would get with these commands

>>> Normal()
Normal(mu=0, sd=1)
>>> Normal('x')
x = Normal(mu=0, sd=1)

And now comes the magic. This will actually work on any class that uses (possibly a subclass of) the type metaclass, but I suppose one will likely want to inherit from some PyMC Model class.

#probably useful to inherit, but not actually necessary for this to work
class MyModel(pm.Model):
    a = Normal(mu=2)

And now check out

>>> model = MyModel()     # not actually necessary to create an instance, but this is the likely workflow 
>>> model.a
a = Normal(mu=2, sd=1)
>>> MyModel.a     # still the same object, not sure if this is is a pro or con ... but can be solved
a = Normal(mu=2, sd=1)

discussion was mainly around removing the yield keyword

Yeah, I definitely support that.

but the main difficulty is how we would get the dict_ dictionary in the @smartass101’s code snippet. I don’t think that object currently exists anywhere in the PyMC3 codebase, so we would have to construct it ourselves… by parsing the AST

Fortunately no, the dict_ namespace is actually populated when the class body is executed and only after that is the metaclass called. So no need to parse the AST, since it’s Python itself doing that.

I’d like to update the post topic to reflect this descriptor easier alternative, but it seems that is not possible anymore.

[smartass101]

Should I perhaps create a separate thread for the simpler descriptor protocol? I cannot change the name of this thread it seems.

I’m curious what you think about this descriptor approach. In particular, I wonder whether my comments above have addressed your reservations @_eigenfoo.

[smartass101]

Has there been any progress on this @junpenglao, @_eigenfoo? I’m just making sure this hasn’t been completely lost to history, because I believe it could bring something very useful to pymc4.

[junpenglao]

We didnt make any progress on this, but if you would like to explore the idea a bit further, we can create a new Github branch that you can submit PR to. How does that sound?

[smartass101]

I’m willing to give it a try, essentially I’d just try to modify the “Distribution” class to have a __set_name__ method. But it would be very helpful if you could point me to some write-up or summary notes of how the current pymc4 implementation instantiates distributions. I’m of course counting on reading the source anyway, but an overview would probably give me a head start.

[junpenglao]

PyMC4 works a bit different so I dont think you can get much information from there (it does not represent RandomVariable as Tensor)

[smartass101]

Yeah, I had a look at the pymc4 design guide notebook and now I see that indeed the co-routines are an integral part of pymc4, not just a syntax workaround (if I understood it correctly) as I thought earlier (I thought a static graph is used like with Theano).

The __set_name__ method can work only in a class-attribute context where the yield generator keyword cannot be used as you do now. So I’d have to later translate the class-attributes (which would kind off lazy-evaluate) as they are defined to an equivalent generator function to be actually used with the current pymc4 design. I don’t think this is impossible, but it would require adding an extra API layer (probably not a big issue if it can be auto-generated based on available tensorflow distribution classes) for representing these “abstract-node-like” distributions.

[smartass101]

I think it would have been simpler to implement with the earlier context-based design that I found described here.
I haven’t be able to easily determine what was the rationale for switching to the coroutine-based design, but I assume there was a good reason and it’s not likely to change back for the sake of a neat auto-naming natural syntax (or is it?).

[smartass101]

I’m thinking whether it might still be worth it implementing this for pymc3 since it is still actively maintained and likely will be for quite some time. Would there be interest in this?

[AlexAndorra]

I guess that could be interesting, provided the implementation adds this novelty optionally – i.e doesn’t break users’ existing code where they specified the var names.
What’s your take @junpenglao?