jax.typing.DTypeLike

jax.typing.DTypeLike = str | type[typing.Any] | numpy.dtype | jax._src.typing.SupportsDType

Represent a PEP 604 union type

E.g. for int | str