Open kazewong opened 4 months ago
Currently, the typing system does not leverage jaxtyping https://docs.kidger.site/jaxtyping/, which helps with static analysis.
For example, in the current typing scheme, a function could look something like
def f(x: Array) -> Array: ...
With jaxtyping, the function should look like:
def f(x: Float[Array, " n_samples"]) -> Float[Array, " n_samples"]: ...
which provides more information to the user as well. I think this is better considering the shape of array is usually very important.
Currently, the typing system does not leverage jaxtyping https://docs.kidger.site/jaxtyping/, which helps with static analysis.
For example, in the current typing scheme, a function could look something like
With jaxtyping, the function should look like:
which provides more information to the user as well. I think this is better considering the shape of array is usually very important.