| # Copyright (c) 2022 Google LLC | |
| # | |
| # Permission is hereby granted, free of charge, to any person obtaining a copy of | |
| # this software and associated documentation files (the "Software"), to deal in | |
| # the Software without restriction, including without limitation the rights to | |
| # use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of | |
| # the Software, and to permit persons to whom the Software is furnished to do so, | |
| # subject to the following conditions: | |
| # | |
| # The above copyright notice and this permission notice shall be included in all | |
| # copies or substantial portions of the Software. | |
| # | |
| # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR | |
| # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS | |
| # FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR | |
| # COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER | |
| # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN | |
| # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. | |
| # Note that `from typing import Annotated; Bool = Annotated` | |
| # does not work with static type checkers. `Annotated` is a typeform rather | |
| # than a type, meaning it cannot be assigned. | |
| from typing import ( | |
| Annotated as BFloat16, # noqa: F401 | |
| Annotated as Bool, # noqa: F401 | |
| Annotated as Complex, # noqa: F401 | |
| Annotated as Complex64, # noqa: F401 | |
| Annotated as Complex128, # noqa: F401 | |
| Annotated as Float, # noqa: F401 | |
| Annotated as Float8e4m3b11fnuz, # noqa: F401 | |
| Annotated as Float8e4m3fn, # noqa: F401 | |
| Annotated as Float8e4m3fnuz, # noqa: F401 | |
| Annotated as Float8e5m2, # noqa: F401 | |
| Annotated as Float8e5m2fnuz, # noqa: F401 | |
| Annotated as Float16, # noqa: F401 | |
| Annotated as Float32, # noqa: F401 | |
| Annotated as Float64, # noqa: F401 | |
| Annotated as Inexact, # noqa: F401 | |
| Annotated as Int, # noqa: F401 | |
| Annotated as Int2, # noqa: F401 | |
| Annotated as Int4, # noqa: F401 | |
| Annotated as Int8, # noqa: F401 | |
| Annotated as Int16, # noqa: F401 | |
| Annotated as Int32, # noqa: F401 | |
| Annotated as Int64, # noqa: F401 | |
| Annotated as Integer, # noqa: F401 | |
| Annotated as Key, # noqa: F401 | |
| Annotated as Num, # noqa: F401 | |
| Annotated as Real, # noqa: F401 | |
| Annotated as Shaped, # noqa: F401 | |
| Annotated as UInt, # noqa: F401 | |
| Annotated as UInt2, # noqa: F401 | |
| Annotated as UInt4, # noqa: F401 | |
| Annotated as UInt8, # noqa: F401 | |
| Annotated as UInt16, # noqa: F401 | |
| Annotated as UInt32, # noqa: F401 | |
| Annotated as UInt64, # noqa: F401 | |
| TYPE_CHECKING, | |
| ) | |
| if not TYPE_CHECKING: | |
| assert False | |
| from jax import ( | |
| Array as PRNGKeyArray, # noqa: F401 | |
| Array as Scalar, # noqa: F401 | |
| ) | |
| from jax.typing import ArrayLike as ScalarLike # noqa: F401 | |