Store results for further uses in the PyTree & jit/vmap context ?

[jecampagne]

Well, the text is certainly too long, sorry. Notably please look at the 2nd snippet after (nb. I keep the first snippet too to see what I was looking for)

I am asking the possibility to store (eventually long) computations in a sort of Work Space in the context of PyTree.
Here is a snippet that I share on Colab in this notebook

Here is the description of the use-case and the snippet.

I have first a base class that implement some commonalities as the tree_flatten/tree_unflatten functions, as well as the storage of two kinds of "parameters tracking" in the following we can forget the gsparams-like parameters allthough I let the code asis.
Then, I have an inherited class that is instantiated with some arguments that will serve to define the Base.params dicttionnary.

For a while let us forget the "ws" parameters

import jax
import jax.numpy as jnp
from jax.tree_util import register_pytree_node_class


class Base:
    def __init__(self, *, gsparams=None, **params):
        self._params = params  # Dictionary containing all traced parameters
        self._gsparams = gsparams  # Non-traced static parameters

    @property
    def gsparams(self):
        """A `GSParams` object that sets various parameters relevant for speed/accuracy trade-offs."""
        return self._gsparams

    @property
    def params(self):
        """A Dictionary object containing all parameters of the internal represention of this object."""
        return self._params

    def __eq__(self, other):
        is_same = self is other
        is_same_class = type(other) is self.__class__
        has_same_trees = self.tree_flatten() == other.tree_flatten()
        return is_same or (is_same_class and has_same_trees)


    def tree_flatten(self):
        """This function flattens the GSObject into a list of children
        nodes that will be traced by JAX and auxiliary static data."""
        # Define the children nodes of the PyTree that need tracing
        children = (self.params,)
        # Define auxiliary static data that doesn’t need to be traced
        aux_data = {"gsparams": self.gsparams}
        return (children, aux_data)

    @classmethod
    def tree_unflatten(cls, aux_data, children):
        """Recreates an instance of the class from flatten representation"""
        return cls(**(children[0]), **aux_data)

@register_pytree_node_class
class A(Base):
  def __init__(
        self,
        beta,
        scale,
        trunc=0.0,
        ws = {},                #### 
        gsparams=None,
  ):

    super().__init__(
                beta=beta,
                scale=scale,
                trunc=trunc,
                ws = {},         #### 
                gsparams=gsparams,
            )
  @property
  def beta(self):
    return self.params["beta"]

  @property
  def trunc(self):
    return self.params["trunc"]

  @property
  def scale(self):
    return self.params["scale"]

  @property
  def _y_trunc(self):
    return  self.scale * self.trunc**self.beta
  
  @property
  def _y_untrunc(self):
    return self.scale**self.beta

  @property
  def _y(self):
    if "_y" in self.params["ws"]:              #### 
      print("use record")
      return  self.params["ws"]["_y"]
    else:
      res = jax.lax.select(self.trunc > 0, self._y_trunc, self._y_untrunc)
      self.params["ws"]["_y"] = res
      return res

  def __hash__(self):
    return hash(
            (
                "A",
                self.beta,
                self.scale,
                self.trunc,
                self.wc,                       #### 
                self.gsparams,
            )
        )

With this snippet it is possible to run

obj1 = A(beta=5.,scale=0.2)
obj1._y

or

obj2 = A(beta=5.,scale=0.2, trunc=2.3)
obj2._y

In a real context, however, the computation of variables such as "_y" can be cumbersome and one ask then why not to use a storage mechanism once the _y is called the first time as it depends only on parameters defined at class construction. Notice that these initialisation cannot be done in the init method as one will get problems with JAX when jit will be advocated.

It is the reason of the "ws" (aka work space) argument and I have put same "####" characters to point where it is used.
If one repeats the above python codes, and for instance print the obj2.params, and repeats the call to obj2._y then one will see after the first call, the message "use record" which signal that the "complicated code" is bypassed to return the result stored in the ws dico.

So far so good, but this is not the end of the story and motivate opening the discussion. In the code development, there are been set some tests, notably a "jitting" and a "vmapping" tests as followed (adapted for the present use case):

identity = jax.jit(lambda x: x)
def test_jit():
  objects = [
        A(beta=5.,scale=0.2),
        A(beta=2.5,scale=0.1, trunc=2.3)
  ]

  def test_eq(self, other):
    return (
            self.scale == other.scale
            and self.gsparams == other.gsparams
        )

  # Check that after jitting the oject is still the same
  return  all([test_eq(identity(o), o) for o in objects])

e = jnp.ones(2)
def duplicate(params):
    return {x: y * e for x, y in params.items()}
def test_vmap():
    objects = [
        A(beta=5.,scale=0.2),
        A(beta=2.5,scale=0.1, trunc=2.3)
    ]

    def test_eq(self, other):
        return (self.scale == jnp.array([other.scale, other.scale])).all() and ( 
            self.beta == jnp.array([other.beta, other.beta])
        ).all()

    # Check that after vmapping the oject is still the same
    return all(
        [test_eq(jax.vmap(A)(**duplicate(o.params)), o) for o in objects]
    )

When the "ws" mechanism is not used at all and we recompute each time "_y" then both test return Ok (aka True).
But when the "ws" mechanism is used then the "vmap" test crash because of the duplicate function. It may exist a workaround to tune this duplicate function. But

My question is more on the use of the "Work Space" mechanism. Do you have any other mechanism that can fit our development philosophy? Thanks.

[jecampagne]

Hi, thinking a bit I see a possible (simple) solution where the "_ws" is just initialized to an empty dictionnary in the init method.
Then, the modified sub-class is the following and all the test_jit/test_vmap are ok, and the successive call to "_y" property is doing its job.

@register_pytree_node_class
class A(Base):
  def __init__(
        self,
        beta,
        scale,
        trunc=0.0,
        gsparams=None,
  ):

    super().__init__(
                beta=beta,
                scale=scale,
                trunc=trunc,
                gsparams=gsparams,
            )
    # Create a workspace where functions can store some precomputed
    # results
    self._ws = {}     #####

  @property
  def beta(self):
    return self.params["beta"]

  @property
  def trunc(self):
    return self.params["trunc"]

  @property
  def scale(self):
    return self.params["scale"]

  @property
  def _y_trunc(self):
    return  self.scale * self.trunc**self.beta
  
  @property
  def _y_untrunc(self):
    return self.scale**self.beta

  @property
  def _y(self):
    if "_y" in self._ws:              #### 
      print("use record")
      return  self._ws["_y"]
    else:
      res = jax.lax.select(self.trunc > 0, self._y_trunc, self._y_untrunc)
      self._ws["_y"] = res
      return res

  def __hash__(self):
    return hash(
            (
                "A",
                self.beta,
                self.scale,
                self.trunc,
                self.gsparams,
            )
        )

But I was wandering if there is not a mechanism to put "property" like function in a cache?

[jecampagne]

no interest?