Type-hinting

[billbrod]

I think this is especially relevant for the jax-casting that happens in #111 , but is a general issue. The question is, how can we effectively use type-hinting for both the static analysis tools and the user. The info we want to communicate is basically: we accept jax arrays, numpy arrays, and pynapple objects. We always return jax arrays or, for those functions decorated with the decorator introduced in #111, pynapple objects (if the input was pynapple)

Posted about this on both the Simons Foundation and US-RSE slacks and got the following advice:

• @overload (docs, PEP) might be useful for this ,but I don't think the info is well-communicated to user.
• should read jax's typing docs for their best practices
  • For the jax/numpy array input leads to jax array output, the following is probably what we want:

    def myfunc(arr: jax.typing.ArrayLike) -> jax.Array: 

• there is a jaxtyping package that allows for annotating array shape and data type.
• someone recommended that using custom classes can capture the "the output will look like the input":

  SpecialType = Union[pynapple.Tsd, jax.numpy.ndarray]
  def myfunc(input: SpecialType) -> SpecialType: