Avoid re-computations of lazy properties after pytree flattening

[mchagneux]

Hi all !

As part of a personal project I have to deal with parameters of multivariate Gaussian distributions a lot, and taking inspiration from NumPyro I want to use some kind of lazy evaluation with some tool like this:

class lazy_property(object):
    r"""
    Used as a decorator for lazy loading of class attributes. This uses a
    non-data descriptor that calls the wrapped method to compute the property on
    first call; thereafter replacing the wrapped method into an instance
    attribute.
    """

    def __init__(self, wrapped):
        self.wrapped = wrapped
        update_wrapper(self, wrapped)

    # This is to prevent warnings from sphinx
    def __call__(self, *args, **kwargs):
        return self.wrapped(*args, **kwargs)

    def __get__(self, instance, obj_type=None):
        if instance is None:
            return self
        value = self.wrapped(instance)
        setattr(instance, self.wrapped.__name__, value)
        return value

which I can use like this for example:

@register_pytree_node_class
@dataclass(repr=True)
class CovParams:

    scale_tril:jnp.ndarray
    cov:jnp.ndarray
    prec:jnp.ndarray


    def __init__(self, scale_tril=None, cov=None):

        if cov is not None:
            self.cov = cov
            self.scale_tril = jnp.linalg.cholesky(self.cov)

        elif scale_tril is not None:
            self.scale_tril = scale_tril
        else:
            raise ValueError(
                "One of `covariance_matrix`, `scale_tril`"
                " must be specified."
            )

    @lazy_property
    def cov(self):
        return jnp.matmul(self.scale_tril, jnp.swapaxes(self.scale_tril, -1, -2))

    def tree_flatten(self):
        return self.scale_tril, None

    @classmethod
    def tree_unflatten(cls, aux_data, params):
        scale_tril = params
        return cls(scale_tril=scale_tril)

The lazy evaluation is practical, but I can't find a way to satisfy both of the following points at the same time:

1. Avoid triggering computation of the lazy properties when the trees are flattened then unflattened (which happens during vmaps, scans, etc)
2. Pass already computed properties to flattened trees and avoid re-trigerring computations for the unflattenned tree.

But maybe under JIT none of this matters ? I am not clear as to what kind of optimizations is done at compile time w.r.t unused computations, e.g. maybe it doesn't matter to compute "cov" if it is not used because the compilation will discard that computation anyway ?

Thanks in advance!

[YouJiacheng]

I think unused computation can trivially be optimized by JIT.
But I think there is some workaround as well:

@register_pytree_node_class
@dataclass(repr=True)
class CovParams:

    scale_tril:jnp.ndarray
    cov:jnp.ndarray
    prec:jnp.ndarray


    def __init__(self, scale_tril=None, cov=None):

        if cov is not None:
            self.cov = cov
            self.scale_tril = jnp.linalg.cholesky(self.cov)

        elif scale_tril is not None:
            self.scale_tril = scale_tril
        else:
            raise ValueError(
                "One of `covariance_matrix`, `scale_tril`"
                " must be specified."
            )

    @lazy_property
    def cov(self):
        return jnp.matmul(self.scale_tril, jnp.swapaxes(self.scale_tril, -1, -2))

    def tree_flatten(self):
        cov = vars(self)['cov']
        if not isinstance(cov, jnp.ndarray):
            cov = None
        return self.scale_tril, cov

    @classmethod
    def tree_unflatten(cls, aux_data, params):
        scale_tril, cov = params
        obj = cls(scale_tril=scale_tril)
        if isinstance(cov, jnp.ndarray):
            setattr(obj, 'cov', cov)
        return obj

[mchagneux]

Ok thanks ! I changed you implementation slightly by using the aux_data field as follows and it works great !

@register_pytree_node_class
class CovParams:

    def __init__(self, chol=None, cov=None, prec=None):

        if cov is not None:
            self.cov = cov
            self.chol = cholesky(cov)

        elif prec is not None:
            self.prec = prec
            self.chol = chol_from_prec(prec)

        elif chol is not None:
            self.chol = chol
            
        else:
            raise ValueError()        

    @lazy_property
    def cov(self):
        return cov_from_chol(self.chol)

    @lazy_property
    def prec(self):
        return inv_from_chol(self.chol)

    @lazy_property
    def log_det(self):
        return log_det_from_chol(self.chol)

    def tree_flatten(self):
        attrs = vars(self)
        children = attrs.values()
        aux_data = attrs.keys()
        return (children, aux_data)

    @classmethod
    def tree_unflatten(cls, aux_data, params):
        obj = cls.__new__(cls)
        for k,v in zip(aux_data, params):
            setattr(obj, k, v)
        return obj

    def __repr__(self):
        return str(vars(self))

Btw I hadn't noticed but using the @DataClass causes all lazy_properties to be evaluated upon creation, so it seems that I shouldn't use it.

One drawback of this though: when some of the fields are evaluated inside a lax.scan, the structure of the carry changes and an error is raised. II don't think there's any workaround to this because the lazy_property is built so as to replace itself with an attribute to avoid any future function call. This is just a minor disadvantage anyway, in the end this just signals that some parameters need to be precomputed before entering the loop and this can be done in a case by case basis. The flexibility of some structure like CovParams is still super useful in other parts of the code.

Thanks again!

[YouJiacheng]

nice code!

[mchagneux]

Thx!

[YouJiacheng]

FYI, after python3.8, there is functools.cached_property do the same thing as your lazy property.

[mchagneux]

I did not know, thank you!