jax.typing.ArrayLike

jax.typing.ArrayLike = jax.Array | numpy.ndarray | numpy.bool | numpy.number | bool | int | float | complex

Represent a PEP 604 union type

E.g. for int | str