-
Notifications
You must be signed in to change notification settings - Fork 51
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
Introduce beartype & fix types #230
Conversation
Going through the notebooks, another issue: |
Yes, JAX will implicitly convert them. So calling |
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>
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
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]) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Sure, let's do that.
gpjax/__init__.py
Outdated
VariationalGaussian, | ||
WhitenedVariationalGaussian) | ||
from jaxtyping import install_import_hook | ||
with install_import_hook("gpjax", "beartype.beartype"): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
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>
@thomaspinder ok I think it's ready to merge now. I hope I didn't accidentally break anything along the way 😅 |
Please check the type of change your PR introduces:
What is the current behavior?
No runtime type checking.
Resolves #202
What is the new behavior?
Runtime type checking.