Fix jax backend sampling with variable names that are not valid identifiers

[aseyboldt]

No description provided.

[ricardoV94]

What are non identifier names? What was failing?

[aseyboldt]

When we have a model like

    with pm.Model() as model:
        a = pm.Data("a/b", shape=2)

or a nested model, then the variables in the generated logp function have different names (they can't be called a/b after all...).
But I was using kwargs with the real variable names as keys to pass the values to the logp function.
So I just switched to positional args, since we know the order.

[ricardoV94]

Maybe call the test invalid_kwarg_names?

How do we do it in PyMC when we have the point func with names like this? Seems like it should also fail because we unpack a dict, or that somehow works?

[aseyboldt]

The shared variables are typically hidden in the pytensor function, and that uses positional args.
And the strange pymc point functions also assume things are in the right order (which has bitten me more than once...)
I called it non-identifier because this happens more or less if name.isidentifier() is false.

[ricardoV94]

I thought the point func just called the underlying function with **state. It's true it doesn't have to interact with shareds, but the name thing should also be an issue there?

[ricardoV94]

Where do we convert it to positional arguments? Inside PyTensor? Because here it's just unpacking the dict: https://github.com/pymc-devs/pymc/blob/main/pymc%2Fpytensorf.py#L617-L624

[ricardoV94]

Anyway I thought the problem was you cant unpack non identifiers in Python

[aseyboldt]

Nope, no limitation anywhere :-)

It is perfectly legal in python (and I think also in jax) to have a function

def foo(**kwargs):
    pass

And then pass it something like {"a/b": 1}. As long as foo takes the kwargs as **kwargs that's allowed.

The problem is that the jax function we generate in dispatch generates something like

def jax_dispatch(a_b):
    pass

and we can't call that function with foo(**{"a/b": 1}), because the name of the variable and the key in the dictionary don't match.

[ricardoV94]

Hmm right I didn't expect those internal variables to ever be called with kwargs

[ricardoV94]

Or for the names to even respect those from the original graph