Check class init arguments: test with checkify?

[jecampagne]

Hi, I have a long standing pb to perform a test on arguments of class init for PyTrees.
Here a use case where Params has two arguments (a,b) and I would like to test that a>=0..

See below that checkify does not detect the second object instantiation error.

Have you an idea to do such test (nb. for the moment we leave to the user the responsability to instantiate with right args)

Thanks.

import jax
jax.__version__
jax.config.update("jax_enable_x64", True)
from jax import random
import jax.numpy as jnp
import numpy as np

from jax.tree_util import register_pytree_node_class
from jax.experimental import checkify

@register_pytree_node_class
class Params:
    def __init__(self, a, b):
        self._a = a
        self._b = b
        jax.debug.print("__init__ Params(a={}, b={})",self._a,self._b)
        # the if code imply a TraceBack  
        #if a<0.0:
        #  raise ValueError("a must be >= 0")

        def check(a):
          checkify.check(a < 0.0, "a needs to be non-negative")
        jittable_check = checkify.checkify(check)
        jittable_check(a)

    def __repr__(self):
        return f"Params(a={self._a}, b={self._b})"

    @property
    def a(self):
        return self._a

    @property
    def b(self):
        return self._b

    def tree_flatten(self):
        children = (self._a, self._b)
        aux_data = None
        jax.debug.print("tree_flatten...")
        return (children, aux_data)

    @classmethod
    def tree_unflatten(cls, aux_data, children):
        jax.debug.print("tree_unflatten...")
        return cls(*children)

and I want to pass a jit test:

jax.clear_caches()
identity = jax.jit(lambda x: x)
def test_jitting():
    objects = [
        Params(a=0.,b=1.1),
        Params(a=-1.0,b=1.1),
    ]

    # Test equality function
    def test_eq(self, other):
        return (
            self.a == other.a
            and self.b == other.b
        )

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

test_jitting()

I get this result w/o any error detected:

__init__ Params(a=0.0, b=1.1)
__init__ Params(a=-1.0, b=1.1)
tree_flatten...
tree_flatten...
tree_flatten...
tree_unflatten...
__init__ Params(a=0.0, b=1.1)
tree_flatten...
tree_flatten...
tree_unflatten...
__init__ Params(a=0.0, b=1.1)
tree_unflatten...
__init__ Params(a=0.0, b=1.1)
tree_flatten...
tree_flatten...
tree_flatten...
tree_flatten...
tree_unflatten...
__init__ Params(a=-1.0, b=1.1)
tree_flatten...
tree_flatten...
tree_unflatten...
__init__ Params(a=-1.0, b=1.1)
tree_unflatten...
__init__ Params(a=-1.0, b=1.1)
tree_flatten...

[jecampagne]

A small up if someone is interested. Thanks

[patrick-kidger]

Yup, I'm aware! Unfortunately JAX doesn't have any other solution to eagerly raise errors!

[jakevdp]

True, this is because raising errors at runtime is not well-supported by XLA. You can hack around that limitation (as in error_if), but it's not an approach we should be recommending to users without large caveats.

[jecampagne]

@jakevdp Ok I note your remark. Do you know if JAX will propose something to check arguments of PyTree is the context I have?

[jakevdp]

We have some ideas about that for the future, but fundamentally, no: there's currently no supported way to raise errors from XLA computations at runtime. It's why, for example, JAX handles out-of-bound indices by clipping rather than raising an error (e.g. JAX Sharp Bits: out of bound indices).

The best way forward is probably to embrace this quirk of JAX/XLA, and write your code accordingly.

[jecampagne]

Hum, "be silent" as the user provides wrong set of arguments is not what I was used to. Anyway, for the current version of our project we have clearly state to the future users the correct domain of validity of the functions.

Now, I was wandering if all the users (and library) already using Equinox (and probably error_if) are aware of the misuses of pure_callback, as en end user I guess no as it is rather tricky, no?

[jecampagne]

So here is the solution thanks to Equinox:

import equinox as eqx

@register_pytree_node_class
class Params:
    def __init__(self, a, b):
        jax.debug.print("__init__ Params(a={}, b={})",a,b)

        a = eqx.error_if(a, a < 0, "a must be >= 0")
        self._a = a
        self._b = b 

    def __repr__(self):
        return f"Params(a={self._a}, b={self._b})"


    @property
    def a(self):
        return self._a

    @property
    def b(self):
        return self._b

    def tree_flatten(self):
        children = (self._a, self._b)
        aux_data = None
        jax.debug.print("tree_flatten...")
        return (children, aux_data)

    @classmethod
    def tree_unflatten(cls, aux_data, children):
        jax.debug.print("tree_unflatten...")
        return cls(*children)

Here is the prints and Traceback

_init__ Params(a=0.0, b=1.1)
__init__ Params(a=-1.0, b=1.1)

---------------------------------------------------------------------------

EquinoxTracetimeError                     Traceback (most recent call last)

[<ipython-input-21-aff0b188db24>](https://localhost:8080/#) in <cell line: 60>()
     58 
     59 
---> 60 test_jitting()

4 frames

[/usr/local/lib/python3.10/dist-packages/equinox/_errors.py](https://localhost:8080/#) in branched_error_if_impl(x, pred, index, msgs, on_error)
    283                     assert type(index) is int
    284                     if on_error == "raise":
--> 285                         raise EquinoxTracetimeError(msgs[index])
    286                     elif on_error == "breakpoint":
    287                         print(msgs[index])

EquinoxTracetimeError: a must be >= 0

[patrick-kidger]

I'm glad this has worked out for you!

And heads-up that you might also like eqx.Module as a more-robust alternative to registering your own pytree class. (Subclassing eqx.Module does this for you, whilst also avoiding several common mistakes -- e.g. your unflatten rule is calling cls(...), which won't work with vmap etc)

[jecampagne]

Sounds interesting too :+
I do not know how to send u a private message to give you a link.

[jecampagne]

BTW:
if I had to plug such condition to trigger an error like that

if self._a == 0. and self._b <= 1:
  raise "for b<=1, then make a>0"

is there a way to do it with equinox error_if?

[patrick-kidger]

Something like
self.a, self.b = eqx.error_if((a, b), (a == 0) && (b <= 1), "some message")
I think!

[jecampagne]

Yes !!! it works :)
For instance

eqx.error_if(None, a==0.0 and  b<=1.0, "Need a>0.0 for b<=1.0")

Sorry our msgs were made at the same times...