Fixed type errors to unblock an internal type annotations refactoring in JAX

google-deepmind / clrs

Apache License 2.0

453 stars 88 forks source link

Fixed type errors to unblock an internal type annotations refactoring in JAX #116

Closed copybara-service[bot] closed 1 year ago

copybara-service[bot] commented 1 year ago

Fixed type errors to unblock an internal type annotations refactoring in JAX

Some JAX internal used Any instead of Array or in their type annotations. https://github.com/google/jax/pull/17760 changed these to alias jax.Array and uncovered type errors fixed here.