JAX array typing confusion

[bwohlberg]

The type name JaxArray is defined in scico.typing, but a number of modules (e.g. scico/linop/_linop.py, scico/linop/_matrix.py, and scico/numpy/util.py) include the import

from jax.interpreters.xla import DeviceArray

and use DeviceArray instead of JaxArray for typing.

It seems most likely that this is just an oversight, but it's also possible that in some cases it's intentional, since JaxArray does not have an identitical definition. This uncertaintly should be resolved, and all instances of importing DeviceArray from jax.interpreters.xla should be replaced by use of an appropriate definition from scico.typing.

[Michael-T-McCann]

For purposes of type annotation, I agree we should make these consistent, with the main goal being to make the documentation look the way we want. But JaxArray seems confusing from the standpoint that you have to look in scico.typing to figure out what it means; it doesn't appear in our documentation or JAX's. I suggest using DeviceArray from jax.numpy (it also exists in scico.numpy, points to the same thing).

Whatever the decision, be careful with find/replacing because there may be places where it's not just a type annotation (thinking of blockarray.py).