Finalise backend-agnosticity of Distributions

[willGraham01]

Creates the Translator class, which expands on the BackendAgnostic ABC. A Translator is aware of the backend_obj (backend object), and also aware of the signature of the required _frontend_methods.

On instantiation, one can pass a sequence of Translation objects - essentially dataclass wrappers for the arguments to convert_signature - to a Translator. For each _frontend_method, the Translator stores the map between the arguments the frontend method takes, and the arguments the corresponding backend method takes. This means that the Translator can be used by the frontend, and will have the expected syntax, but also takes care of the necessary mapping of the arguments provided by the frontend to those that the backend takes. There is a simple example of this in action in the tests for the Translator class, however the tests for different backends demonstrate how this functionality is envisioned to be used.

The Distribution class now inherits from Translator. The NativeDistribution class has also been introduced, for us to use as the "base" for any standard distributions that we want to provide via our jax-default backend. This class simply defaults all the "translations" to the identity map (IE, the arguments and method names of the backend are already what we expect of the frontend).

Users should use the Distribution class when they want to use a backend that is different to our "standard" backend (jax). If they envision using multiple distributions from the same backend, they can create a derived class to reduce the number of times they need to provide the mapping information, for example

from numpyro.distributions.continuous import MultivariateNormal

from causalprog.distributions.base import Distribution
from causalprog.backend.translation import Translation

NUMPYRO_TRANSLATIONS = (
    Translation(backend_method="sample", frontend_method="sample", param_map = {"seed": "rng_key"},
    ...
)

class NumpyroDistribution(Distribution):

    def __init__(self, *, backend, label: str) -> None:
        super().__init__(*NUMPYRO_TRANSLATIONS, backend=backend, label=label)      

# Specialized classes are further possible,
# if the user so desires.

class NumpyroNormal(NumpyroDistribution):

    def __init__(self, mean, cov, *, label: str) -> None:
        super().__init__(backend=MultivariateNormal(mean, cov), label=label)

import jax.numpy as jnp

mean, cov = jnp.array(...), jnp.array(...)
normal = NumpyroNormal(mean, cov, label="Numpyro Normal")
normal.sample(...) # Works with frontend syntax, but calls `numpyro` functionality.

Other Changes

• convert_signature now returns the function that maps the frontend arguments to their backend counterparts, rather than the function that does this and then evaluates the backend function. This is so that we can recycle the static identity map method of Translator.

[willGraham01]

Self-review with some typos

[willGraham01]

Closing this as it's outdated, especially in light of our decision - for now - to commit to numpyro for our distribution needs.