FEAT: define pytrees for bilby types

[ColmTalbot]

This PR adds JAX pytree definitions for most classes in Bilby including:

• all (I think) priors
• all likelihoods
• the GW interferometer and waveform generator classes

Defining pytrees allows these classes to be used as inputs to JIT-compiled (or other transformations) functions and can dramatically reduce compilation/execution time.

As an example, I've used this in a case where I'm analyzing lots of events in series.
I can then JIT-compile an outer function that takes a likelihood as input and everything works smoothly.

The basic breakdown of pytree definitions relies on breaking the constituent parts of the class into leaves and auxillary data.
Any change in the auxillary data causes a recompilation, whereas leaves are traced through the functions allowing for them to be changed without recompilation.
As an example, the shape of an array would be auxillary data, but the values themselves a leaf/leaves.

TODO:

• more testing, especially of likelihoods

[mj-will]

On the call, we suggested adding a test to ensure the likelihood doesn't recompile when it shouldn't.

[GregoryAshton]

One top-level thought. Is it possible to add a README or something to the top-level compat? Right now, the functionality of this is quite hidden, while exposing it fits into the broader theme of user-friendliness. I often poke around a project to see what it has an "compat" isn't enough for me right now to understand that.

[GregoryAshton]

The behaviour of PowerLaw is now different with alpha=-1, and when the array values are small, this could cause differences. Can we add some tests like this to test for that?

def test_powerlaw_alpha_minus_one_rescale(self):
    """Test rescale with alpha=-1 matches LogUniform behavior."""
    prior = PowerLaw(alpha=-1, minimum=1, maximum=10)
    domain = self.xp.linspace(0, 1, 50)
    rescaled = prior.rescale(domain)
    # Verify no NaN/inf in rescaled values
    self.assertTrue(self.xp.all(self.xp.isfinite(rescaled)))
    # Verify inverse CDF property
    cdf_vals = prior.cdf(rescaled)
    self.assertTrue(self.xp.allclose(domain, cdf_vals, atol=1e-9))

def test_powerlaw_alpha_minus_one_prob_and_ln_prob(self):
    """Test prob/ln_prob with alpha=-1 (critical edge case)."""
    prior = PowerLaw(alpha=-1, minimum=1, maximum=10)
    val = self.xp.array([1.0, 5.0, 10.0])
    prob = prior.prob(val)
    ln_prob = prior.ln_prob(val)
    # Verify consistency: ln_prob = log(prob)
    self.assertTrue(self.xp.allclose(ln_prob, self.xp.log(prob), atol=1e-9))
    # Verify no NaN/inf leakage
    self.assertTrue(self.xp.all(self.xp.isfinite(prob)))
    self.assertTrue(self.xp.all(self.xp.isfinite(ln_prob)))

def test_powerlaw_numeric_stability_small_values(self):
    """Test numerical stability near minimum with various alpha values."""
    for alpha in [-1, -0.5, 0, 1, 2]:
        prior = PowerLaw(alpha=alpha, minimum=1e-10, maximum=1)
        val = self.xp.array([1e-10, 1e-5, 1e-1])
        prob = prior.prob(val)
        ln_prob = prior.ln_prob(val)
        # No NaN/inf
        self.assertTrue(self.xp.all(self.xp.isfinite(prob[val >= prior.minimum])))
        self.assertTrue(self.xp.all(self.xp.isfinite(ln_prob[val >= prior.minimum])))

[GregoryAshton]

We discussed on the call adding these back in

[GregoryAshton]

@ColmTalbot gave a comment on the call as to why this is required. Could that be added as a comment as this otherwise looks like a typo