File size: 2,862 Bytes
9c6594c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 |
# 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
|