Scalarlike now just goes by its name in equinox ecosystem documentation

[johannahaffner]

As suggested in patrick-kidger/optimistix#109.

I built the optimistix documentation with jaxtyping from this branch, and it works splendidly.

[patrick-kidger]

Actually, this raises a question I should have spotted earlier -- ScalarLike is a fairly opaque type annotation to see in the documentation. The fact that jaxtyping.ScalarLike exists is not super clear, and it's pretty niche!

FWIW the convention I've been using elsewhere is to strip jaxtyping shape/dtypes in documentation (mostly as I haven't wanted to assume familiarity with jaxtyping syntax for those reading the docs):

https://github.com/patrick-kidger/equinox/blob/550967c2f12f65d14489153650680af3cb6030c6/mkdocs.yml#L82

In which case perhaps we want this to be rendered as ArrayLike?

[johannahaffner]

In optimistix, we use it to type values that should be ArrayLike with shape (). So that extra bit of information would get stripped along with the rest. I think it should be clear that e.g. a learning_rate is a single scalar value.

[patrick-kidger]

I agree, it should be clear! Which is why I'm thinking of stripping that information here for (a) readability for those unfamiliar with jaxtyping and (b) consistency.

As an alternative, we could consider turning on jaxtyping annotations in documentation for the whole ecosystem? The library has quite possibly achieved critical mass at this point, that a large enough fraction of readers will know what it means?

[johannahaffner]

As an alternative, we could consider turning on jaxtyping annotations in documentation for the whole ecosystem? The library has quite possibly achieved critical mass at this point, that a large enough fraction of readers will know what it means?

I would prefer that, I find jaxtyping's expressiveness very useful!

To get back to the previous example - for a learning rate it might not be too confusing to use ArrayLike, since we can assume that people know that this is a single scalar (small) value. But take lambda_0 in IndirectDampedNewtonDescent - annotating this as ArrayLike could lead users to think that this could be a vector quantity (here).

[patrick-kidger]

In that case at minimum let's replace this here to just return Shaped[jaxtyping.ArrayLike, ""] rather than Shaped[jax.typing.ArrayLike, ""] so that we redispatch to the doc-overloaded version of ArrayLike that we have! That'll then give us consistent behaviour regardless of the choice of using jaxtyping annotations in docs.

[johannahaffner]

In that case at minimum let's replace this here to just return Shaped[jaxtyping.ArrayLike, ""] rather than Shaped[jax.typing.ArrayLike, ""] so that we redispatch to the doc-overloaded version of ArrayLike that we have! That'll then give us consistent behaviour regardless of the choice of using jaxtyping annotations in docs.

"here" being where? jaxtyping.ArrayLike is not available in __init__ of jaxtyping itself, and ArrayLike is not available where ScalarLike is defined in our current branched setup. Do you mean in optimistix?

More specifically, I find the Array{ , Like} as part of the name describing a scalar quantity mildly unintuitive where users might also pass a float, since it makes me think that things should be a tuple, or a list. Personally I find ScalarLike a very descriptive name :)