Docs custom rules

[vadmbertr]

Hello,

I ran into a few problems trying to run the custom rules tutorial:

1

Encountered TypeError: '<' not supported between instances of 'Dimension' and 'Dimension' as JAX (v0.4.34 in my env.tries to sort pytrees dictionary keys when flattening. Added methods to allow comparing Dimension using their name attribute.

2

Noticed that

@quax.register(jax.lax.integer_pow_p)
def _(x: Unitful, *, y: int):
    units = {k: v * y for k, v in x.units.items()}
    return Unitful(x.array, units)

does not elevate x.array to the power of y. (incidentally fixes an inconsistent output)

3

Noticed that from jaxtyping import ArrayLike # https://github.com/patrick-kidger/quax does not use the correct URL.

Are you ok with the proposed changes?

Vadim

[patrick-kidger]

These fixes all look good to me! Thank you for the pull request - merged :)