Introduce beartype & fix types #230

st-- · 2023-04-13T14:55:32Z

Please check the type of change your PR introduces:

Bugfixes identified by added type checking
Some backwards-incompatible changes to fix behaviour
Code style update (formatting, renaming)

What is the current behavior?

No runtime type checking.

Resolves #202

What is the new behavior?

Runtime type checking.

…not ideal :S

st-- · 2023-04-25T06:49:40Z

Going through the notebooks, another issue: Float[Array, "..."] specifically matches JAX-arrays, but not numpy arrays. How do JAX-functions handle numpy-arrays? Do they implicitly convert them? Should we therefore relax type hints on all user-facing parts (not just Dataset but also e.g. predict()) to allow numpy arrays?

thomaspinder · 2023-04-25T07:05:53Z

Going through the notebooks, another issue: Float[Array, "..."] specifically matches JAX-arrays, but not numpy arrays. How do JAX-functions handle numpy-arrays? Do they implicitly convert them? Should we therefore relax type hints on all user-facing parts (not just Dataset but also e.g. predict()) to allow numpy arrays?

Yes, JAX will implicitly convert them. So calling posterior(np_array, D) will produce the same result as posterior(jnp.asarray(np_array), D). I think it'd place less friction on the user if we relax the type hint to accomodate Numpy arrays.

gpjax/dataset.py

examples/graph_kernels.pct.py

Co-authored-by: st-- <st--@users.noreply.github.com> Signed-off-by: Thomas Pinder <tompinder@live.co.uk>

Allow for integer responses Co-authored-by: st-- <st--@users.noreply.github.com> Signed-off-by: Thomas Pinder <tompinder@live.co.uk>

thomaspinder

All looks good to me. There's a small issue with imports, but I'm happy to merge as I've setup a pre-commit action on docs_update that will fix this. I will rebase that branch and correct.

If you're happy with the proposed approach of pushing beartype onto the user, then we can merge this.

gpjax/kernels/base.py

gpjax/kernels/nonstationary/linear.py

thomaspinder · 2023-04-25T20:17:05Z

gpjax/linops/linear_operator.py

+DTypes = Union[Type[jnp.float32], Type[jnp.float64], Type[jnp.int32], Type[jnp.int64]]
 ShapeT = TypeVar("ShapeT", bound=NestedT[Tuple[int, ...]])
-DTypeT = TypeVar("DTypeT", bound=NestedT[jnp.dtype])
+DTypeT = TypeVar("DTypeT", bound=NestedT[DTypes])


Sure, let's do that.

thomaspinder · 2023-04-25T20:27:02Z

gpjax/__init__.py

-                                   VariationalGaussian,
-                                   WhitenedVariationalGaussian)
+from jaxtyping import install_import_hook
+with install_import_hook("gpjax", "beartype.beartype"):


That's a good way to position it. Let's push it onto users then. If it becomes an issue, we can always walk it back with little major disruption.

…into st/beartype

gpjax/dataset.py

Signed-off-by: st-- <st--@users.noreply.github.com>

Co-authored-by: st-- <st--@users.noreply.github.com> Signed-off-by: Thomas Pinder <tompinder@live.co.uk>

st-- · 2023-04-26T19:54:58Z

@thomaspinder ok I think it's ready to merge now. I hope I didn't accidentally break anything along the way 😅

st-- added 29 commits April 13, 2023 12:55

add beartype dependency

15b808f

from typing import -> from beartype.typing import

0eaec30

jaxtyping import_hook for @jaxtyped @beartype everywhere

781bc38

fix Type[] of class-as-argument

67e315f

fix KeyArray type hint (should probably move into jaxutils though)

e77732b

fix return value of slice_input when active_dims is None

5175fa7

fix return value of squared_distance

c548d73

fix return type of recursive_bijectors

079936f

fix slice_input type annotations

1005212

new KernelCallable type to fix kernel_fn annotations

0c57a4e

fix kernel __call__ annotation

777b042

fix KeyArray type hint

df88c69

beartype does not like forward references; replaced with string types

711fcac

linops other type hint fixes

1c95fa5

fix KeyArray

1b915cf

abstractions.py some type fixes

c3e48c1

fix GaussianDistribution.log_prob return type

4cc9db1

fix depreciations & warnings

4c4a572

Merge branch 'st/fix_depreciations' into st/beartype

c512195

fix scalar array types

9516fdf

introduce ScalarBool, ScalarInt for jitted calls in abstractions

2d45258

relax LinearOperator's solve() types (can be both matrix or vector), …

ecfde64

…not ideal :S

remove _stop_grad type hints, not sure what they should be

066a5d7

found some more

ff68725

Merge branch 'st/fix_depreciations' into st/beartype

4716773

float -> ScalarFloat fixes

5c90755

linops log_det type fixes

c692568

some more linops type fixes

ff5211a

Merge remote-tracking branch 'upstream/v0.6' into st/beartype

f3cc17e

st-- changed the base branch from main to v0.6 April 13, 2023 15:43

st-- added 5 commits April 25, 2023 18:15

our own Array type that accepts both JAX and Numpy arrays

edaf2cb

some Float -> Num relaxations for graph kernel...

bb4d51f

ScalarFloat for GraphKernel hyperparams

ea1537d

fix type hints to what happens (even if it seems wrong)

495c649

type relaxation for deep_kernels.pct.py

dca6799

st-- mentioned this pull request Apr 25, 2023

bug: GraphKernel does not conform to AbstractKernel interface #239

Open

st-- requested a review from thomaspinder April 25, 2023 15:57

st-- commented Apr 25, 2023

View reviewed changes

gpjax/dataset.py Outdated Show resolved Hide resolved

st-- commented Apr 25, 2023

View reviewed changes

gpjax/dataset.py Outdated Show resolved Hide resolved

st-- commented Apr 25, 2023

View reviewed changes

examples/graph_kernels.pct.py Outdated Show resolved Hide resolved

st-- and others added 4 commits April 25, 2023 19:18

some more minor consistency fixes

f26d52e

bugfix

23efd03

Update examples/graph_kernels.pct.py

1077efa

Co-authored-by: st-- <st--@users.noreply.github.com> Signed-off-by: Thomas Pinder <tompinder@live.co.uk>

Update gpjax/dataset.py

c292c7e

Allow for integer responses Co-authored-by: st-- <st--@users.noreply.github.com> Signed-off-by: Thomas Pinder <tompinder@live.co.uk>

thomaspinder approved these changes Apr 25, 2023

View reviewed changes

st-- added 4 commits April 26, 2023 16:34

jaxtyping import hook for notebooks

7e16487

conftest.py to apply jaxtyping import hook before loading tests

6620802

remove import hook from gpjax/__init__

bc70fdf

Merge branch 'v0.6' of https://github.com/JaxGaussianProcesses/GPJax …

249c9d0

…into st/beartype

st-- commented Apr 26, 2023

View reviewed changes

gpjax/dataset.py Outdated Show resolved Hide resolved

st-- commented Apr 26, 2023

View reviewed changes

gpjax/dataset.py Outdated Show resolved Hide resolved

st-- and others added 3 commits April 26, 2023 17:00

Update gpjax/dataset.py

c4a8af8

Signed-off-by: st-- <st--@users.noreply.github.com>

Update gpjax/dataset.py

62abd1a

Co-authored-by: st-- <st--@users.noreply.github.com> Signed-off-by: Thomas Pinder <tompinder@live.co.uk>

fix tests of shape checks now that we have beartype

42e21fb

thomaspinder merged commit 75f9ddd into JaxGaussianProcesses:v0.6 Apr 27, 2023

st-- mentioned this pull request Apr 27, 2023

feat: Use beartyping #202

Closed

st-- deleted the st/beartype branch November 22, 2023 21:38

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Introduce beartype & fix types #230

Introduce beartype & fix types #230

st-- commented Apr 13, 2023 •

edited

Loading

st-- commented Apr 25, 2023

thomaspinder commented Apr 25, 2023

thomaspinder left a comment

thomaspinder Apr 25, 2023

thomaspinder Apr 25, 2023

st-- commented Apr 26, 2023

Introduce beartype & fix types #230

Introduce beartype & fix types #230

Conversation

st-- commented Apr 13, 2023 • edited Loading

What is the current behavior?

What is the new behavior?

st-- commented Apr 25, 2023

thomaspinder commented Apr 25, 2023

thomaspinder left a comment

Choose a reason for hiding this comment

thomaspinder Apr 25, 2023

Choose a reason for hiding this comment

thomaspinder Apr 25, 2023

Choose a reason for hiding this comment

st-- commented Apr 26, 2023

st-- commented Apr 13, 2023 •

edited

Loading