Fix _DistributionTerm.expand return-type annotation

[datvo06]

Closes #666.

Per @eb8680's feedback — sticking to NumPyro's Distribution.expand signature (-> Distribution) rather than monkey-patching it to -> Self. Reverted the typing.Self annotation attempt and the merge of #613. Back to registering a defdata case for dist.Distribution (per @jfeser's suggestion in the same thread).

Three small changes in effectful/handlers/numpyro.py:

1. _DistributionTerm.expand return annotation: jax.Array → dist.Distribution. Was routing .expand([J]) through _ArrayTerm, breaking .to_event(1) chaining.

2. _DistributionTerm._is_eager recurses into _DistributionTerm-valued args (rather than treating them as non-eager via is_eager_array). Otherwise a fully-eager D.to_event(k) would be tagged non-eager and every downstream property/sample method raised.

3. defdata.register(dist.Distribution) mapping to _DistributionMethodTerm. Without it, dist.Distribution falls through to _CallableTerm (since Distribution defines __call__). The class overrides _pos_base_dist to materialise by getattr(base, self._op.__name__)(*args[1:], **kwargs), so every inherited @defop (batch_shape, event_shape, support, log_prob, sample, …) resolves through the standard machinery against the real NumPyro distribution.

Three regression tests (each verified to fail on master):

• narrow chain construction (the exact AttributeError from _DistributionTerm.expand and .to_event return _ArrayTerm instead of Distribution, breaking standard NumPyro vectorization idiom #666)
• equational shape/support laws on .expand / .to_event
• end-to-end MCMC over Normal(mu_term, 1.0).expand([3]).to_event(1) with mu_term bound via effectful.ops.semantics.handler

Beta(...).to_event(k) / Dirichlet(...).to_event(k) parametrised cases stay xfail — the basic chain works, but composing .expand_by on top with indexed (named-dim) arrays hits a deeper indexed-dim integration matter, out of scope for this PR.

Lint clean (ruff / format / mypy).

[datvo06]

Not ready yet. Reproduction in #666 still fails, now as (_CallableTerm instead of _ArrayTerm previously). Also havent added the reproduction test.

[datvo06]

Fixed expand().to_event(), but for full MCMC examples in #666, I'd need to update it, since NumPyro'd eventually need mu() actual value.

[datvo06]

Ready now!

[jfeser]

This test doesn't expand a symbolic distribution, because mu is always concrete. Did it fail before?

[datvo06]

Working on merging #613 and find a cleaner solution.

[datvo06]

@jfeser I ran into this problem of using unify on an inherited type:

  import typing                                                                                                                                                                                                                                                       
  from effectful.internals.unification import unify                                                                                                                                                                                                                   
  
  T = typing.TypeVar("T")                                                                                                                                                                                                                                             
                                                                                                                                                                                                                                                                    
  class Foo: pass
  class Bar(Foo): pass

  assert unify(T, Bar) == {T: Bar}      # FAILS: gets {T: Foo}          

So whenever you unify an inherited type with T, it resolves to the parent class. Should we default to principal type or widening in unifying/type inferring?
This is blocking the Self-based routing, because then Normal(mu()).expand(s) would result in dist.Distribution instead of NormalTerm via Self -> dist.Normal anddispatch falls back to _CallableTerm.
Which raises another error.

[jfeser]

That's definitely a bug.

[eb8680]

I am not sure about the premise of this PR. The latest proposed typing rule is not consistent with the signature of numpyro.distributions.Distribution.expand, which only promises to return a Distribution, not something of the same type as self:

def expand(self, batch_shape: Sequence[int]) -> Distribution: ...

I'm missing context on what you're trying to accomplish but I think it's unwise in general to attempt to monkey-patch typing behavior of upstream libraries within effectful like this and I'd strongly recommend sticking to the NumPyro signature for expand.

[jfeser]

Sorry, that's my mistake. In that case, we might register _DistributionTerm as a defdata case for dist.Distribution.

[datvo06]

Oh got it. Will revert back to _DistributionTerm defdata

The context was that I was attempting to factor out the interfacing between effectful symbols and numpyro here. So then you could do something like this:

import jax.numpy as jnp
import jax.random as jr

from effectful.handlers.numpyro import Normal, HalfNormal
from effectful_mcmc import sample, MCMC, NUTS

def linear_regression(x, y):
    alpha = sample(Normal(0.0, 10.0))
    beta  = sample(Normal(0.0, 10.0))
    sigma = sample(HalfNormal(2.0))
    sample(Normal(alpha + beta * x, sigma), obs=y)
    return alpha, beta, sigma                      # return per-site effects

mcmc = MCMC(NUTS(linear_regression), num_warmup=500, num_samples=1000)
mcmc.run(jr.PRNGKey(0), x_data, y_data)

# Posteriors are keyed by the Operation behind each sample site;
# `term.op` recovers it. No string names to keep in sync.
alpha, beta, sigma = mcmc.model_return_value
samples = mcmc.get_samples()
print(float(samples[alpha.op].mean()),
      float(samples[beta.op].mean()),
      float(samples[sigma.op].mean()))

Which is a bit more like normal Python. I was thinking something like this could be factored out of RoboTL.