Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Coherent policy needed on array types #181

Open
bwohlberg opened this issue Jan 18, 2022 · 5 comments
Open

Coherent policy needed on array types #181

bwohlberg opened this issue Jan 18, 2022 · 5 comments
Assignees
Labels
discussion required Some discussion necessary to decide how to address this issue improvement Improvement of existing code, including addressing of omissions or inconsistencies

Comments

@bwohlberg
Copy link
Collaborator

This issue is related to both typing and explicit array conversion and creation (the latter form being closely related to the subject of #93). At the moment we have a somewhat random mix of Array and JaxArray type specifications within our code. A coherent policy is needed as a guide to addressing the existing random usage, and ensuring that it doesn't persist into new code.

The simplest policy is obviously to simply require that all scico functions and methods take and return JaxArray types. Is this both feasible and desirable? If not, what are the counter-examples?

@bwohlberg bwohlberg added the discussion required Some discussion necessary to decide how to address this issue label Jan 18, 2022
@lukepfister
Copy link
Contributor

Make it all JaxArray and be done with it. But note that will really mess with static typing, as jax arrays are really aliased to Any.

Note also that using np.ndarray may not be the "correct" way to annotate a function that takes an ndarray input, cf the numpy typing module

See also the jax policy on type checking: jax-ml/jax#8224 (tl, dr; make everything alias to Any)

@bwohlberg bwohlberg added the improvement Improvement of existing code, including addressing of omissions or inconsistencies label Jan 21, 2022
@tbalke
Copy link
Contributor

tbalke commented Jan 24, 2022

First, let's have a type, i.e. ScicoArray = Union[BlockArray, DeviceArray] or whatever we want to call it.

I think there are arrays that don't make sense to be ScicoArrays. For example, the angle list of a projector. It is rather an array-shaped parameter, not any different than the pixel-spacing of the projector (other than that is an array). If we choose to make this an ScicoArray then we may as well make the pixel-spacing a zero dimensional ScicoArray, but that would obviously not be sensible. These cases can be easily identified and should only be numpy arrays or python scalars.

Then there are operators, like f(x) = a B x where a is a scaler, B is a matrix, x is a vector.
In my opinion B and x should be ScicoArrays but a should not be required to be a ScicoArray. In any case operators should always return ScicoArrays (related to #93).

@bwohlberg
Copy link
Collaborator Author

A good part of this issue has been addressed in #410, but the question of whether all scico functions should take and return jax arrays (as opposed to numpy arrays) is still open.

@bwohlberg
Copy link
Collaborator Author

It's worth noting the availability of jax.typing.ArrayLike

for any value that is safe to implicitly cast to a JAX array; this includes jax.Array, numpy.ndarray, as well as Python builtin numeric values (e.g. int, float, etc.) and numpy scalar values (e.g. numpy.int32, numpy.flota64, etc.)

Also, see section "JAX Typing Best Practices" in jax.typing.

@bwohlberg bwohlberg self-assigned this May 22, 2023
@bwohlberg
Copy link
Collaborator Author

Suggestion: pre- and post-processing functions should take numpy or jax arrays, while functions expected to be used within optimization problems etc. should take jax arrays.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
discussion required Some discussion necessary to decide how to address this issue improvement Improvement of existing code, including addressing of omissions or inconsistencies
Projects
None yet
Development

No branches or pull requests

3 participants