[Docos] Add a new FAQ entry on preventing beartype.claw from type-checking "bad boy" callables and classes

[SalamanderXing]

Hi! I just tried the new amazing feature beartype_this_package() and I'm running into an issue, here's an example:

    params = model.init(
             ^^^^^^^^^^^
File "<@beartype(graph_diffusion.models.gt_digress.main.GraphTransformer.__init__) at 0x7f0014350f40>", line 103, in __init__
beartype.roar.BeartypeCallHintParamViolation: Method graph_diffusion.models.gt_digress.main.GraphTransformer.__init__() parameter parent=<flax.core.scope.Scope object at 0x7f000cb0e350> violates type hint typing.Union[typing.Type[flax.linen.module.Module], typing.Type[flax.core.scope.Scope], typing.Type[flax.linen.module._Sentinel], NoneType] under non-default configuration BeartypeConf(claw_is_pep526=True, is_color=None, is_debug=False, is_pep484_tower=False, strategy=<BeartypeStrategy.O1: 2>warning_cls_on_decorator_exception=<class 'beartype.roar.BeartypeClawDecorWarning'>, ), as <class "flax.core.scope.Scope"> <flax.core.scope.Scope object at 0x7f000cb0e350> not type or <class "builtins.NoneType">.

So from my understanding, it looks like the authors of the upstream package Flax package have annotated a parameter as typing.Type[TheType] instead of TheType directly and that causes everything to fail.

Is there a way to pause beartype for that specific call?

[leycec]

Aww. Thanks so much for such kind words that both massage my aching soul and stroke my ego gently. Let's see if @beartype can help set this typing conundrum right for the good of AI in Amsterdam. 🤖 + ☕ = win!

Now, If I'm reading that type violation right: ^{...which I'm probably not, because the Canadian heat wave has fried my neocortex like an overcooked pancake on the outdoor skillet.}

• Your package that you have augmented with beartype_this_package() is named graph_diffusion. Is that right? Google shows nothing, suggesting your package has yet to be open-sourced into the wild – perhaps because you are still conducting clandestine research or quietly making clandestine money off this thing. I applaud and approve of both.
• Your graph_diffusion.models.gt_digress.main namespace ^{...module? function? I am clueless, which is par for the course around here.} defines a class named GraphTransformer.
• Your GraphTransformer.__init__() method accepts a parameter named parent annotated by the type hint typing.Union[typing.Type[flax.linen.module.Module], typing.Type[flax.core.scope.Scope], typing.Type[flax.linen.module._Sentinel], NoneType]. The signature for this method resembles:

class GraphTransformer(...):
    def __init__(self, parent: typing.Union[
        typing.Type[flax.linen.module.Module],
        typing.Type[flax.core.scope.Scope],
        typing.Type[flax.linen.module._Sentinel],
        NoneType
    ], ...) -> None: ...

Is that all right? I suppose what I'm trying (but failing) to obliquely suggest is that the beartype_this_package() import hook should only apply to your graph_diffusion package; beartype_this_package() definitely should not apply to Google Flax, because neither you or I are Google. I know I'm not Google and I'm pretty sure you're not, either. ^{...but what if one of us is Google? mind blown.}

Therefore, musn't it be the case that your GraphTransformer.__init__() method (rather than something in Flax or any other upstream package) is incorrectly annotated? In short, shouldn't that method instead resemble:

class GraphTransformer(...):
    def __init__(self, parent: typing.Union[
        flax.linen.module.Module,
        flax.core.scope.Scope,
        flax.linen.module._Sentinel,
        None
    ], ...) -> None: ...

Admittedly, I have no idea what's going on here. Curse you yet again, vicious Ontario heat wave.

For example, why does that type hint reference the private Flax type flax.linen.module._Sentinel? That's a pretty weird violation of privacy encapsulation – which substantiates your claim that this is, indeed, an incorrectly typed Flax thing. But... how could that be? Perhaps beartype_this_package() is doing horrible things. Perhaps I should go to bed before I begin contemplating those horrible things. 😪

[SalamanderXing]

Thank you for your answer! And sorry for not explaining better the misleading stack trace.
Here's the source code for the google code: https://github.com/google/flax/blob/9ba7cc0cc54a5eb38df637892358cc9e1006e225/flax/linen/module.py#L737C7-L737C17
Although it looks from the stack trace that the __init__ method of my GraphTransformer class, that is not the case. That method is inherited from the flax.linen.Module class (linked above) and I have not explicitly defined a parameter called "parent".

class GraphTransformer(nn.Module):
    n_layers: int
    hidden_mlp_dims: Dims = Dims(x=128, e=128, y=128)
    hidden_dims: HiddenDims = HiddenDims(
        x=256, e=64, y=64, n_head=8, dim_ffx=256, dim_ffe=128, dim_ffy=2048
    )

    @classmethod
    def initialize(
        cls,
        key: Key,
        in_node_features: int,
        in_edge_features: int,
        number_of_nodes: int,
        num_layers: int,
    ):
        ...

    @nn.compact
    def __call__(self, g: gd.OneHotGraph, y, deterministic: bool):
        ...

[leycec]

Ah-ha! Fantastic. Thanks so much for the minimal-reproducible example (MRE). You are totally right about everything. Indeed, this is an upstream Google Flax typing bug. Woooooooah. Never expected Google of all tech giants and Flax of all tech giant ML frameworks to drop the QA ball, but... we're all AI-augmented humans here, right? Mistakes do happen and then @beartype raises its claw in anger against those mistakes.

Would you mind submitting an upstream issue against Flax's issue tracker? Interestingly, this further substantiates the need for beartype.claw-fueled hybrid runtime static type-checking. After all, Google (probably) already statically type-checks Flax against their own in-house pytype static type-checker – which mistakenly thought everything was just hunky-dory. But Flax isn't just hunky-dory, is it, Google? </sigh>

@leycec: Does That Neckbeard Ever Shuddap?

The simple answer is: "No." Since you've already done so much, allow me to now temporarily resolve this issue for you on your end until Google themselves permanently resolve this for everybody.

@beartype supports both the standard @typing.no_type_check decorator and the @beartype-specific beartype.BeartypeStrategy.O0 no-time strategy. Both effectively blacklist a decorated callable or class from runtime type-checking by @beartype. Sadly, since this specific issue pertains to a superclass method outside your control, you kinda have no recourse but to temporarily blacklist your entire GraphTransformer class.

In other words, choose to do either of these two things:

# Import the requisite machinery.
from beartype import beartype, BeartypeConf, BeartypeStrategy  # <-- this works
#from typing import no_type_check  # <------------------------------ this works, too!

# Dynamically create a new @nobeartype decorator disabling type-checking.
nobeartype = beartype(conf=BeartypeConf(strategy=BeartypeStrategy.O0))

# Avoid type-checking *ANY* methods or attributes of this class.
@nobeartype  # <------ this works
#@no_type_check  # <-- this works, too!
class GraphTransformer(nn.Module):
    ...

Since this is kinda unintuitive, I've:

• Expanded our existing "...as Noop" documentation on the subject.
• Authored a new FAQ entry "How do I NOT type-check something?".

Thanks so much for flying the friendly QA skies with @beartype. We hope you have your enjoyed your stay on Bear Bro Airlines an– ohgodsohgods the wing's on fire!!!!!!! ✈️ 💥

[SalamanderXing]

Thank you so much @leycec ! Issue is here: google/flax#3224