Commit Graph
14 Commits
Author SHA1 Message Date
Patrick Kidger 849b15db3b Added better docs; added PRNGKeys. Packaging with pyproject.toml.
Also added support for matching dtypes using regexes, and for nesting e.g. Shaped[Float[Array, "dim1 dim2"], "dim3"].
2023-05-10 15:27:25 -07:00
Patrick Kidger c92b0d0ab1 Some improvements (#77)
* Various improvements.

- Added support for functions in symbolic dimensions, e.g. "min(foo,bar)", which were previously disallowed due to the presence of a comma. (#51)
- Added support for adding ignored names to dimensions, e.g. "cols=4". (#76)

* Now works with Python 3.10 A | B union types.
2023-04-13 19:18:55 +01:00
Patrick Kidger e03c1c329e We now have Float[np.ndarray, ...] <: np.ndarray. Added basic torch tests. (#68)
This required quite a lot of refactoring! JAX supports virtual subclass registration (its metaclass is ABCMeta) but NumPy does not, so we have to actually subclass `np.ndarray`.
Simple stuff like __base__ hacking fails due to deallocator conflicts.
2023-03-04 17:29:04 +00:00
Patrick Kidger 5600a1aac8 Fixed cloudpickle breaking, mark 2 2023-03-02 17:37:53 -08:00
Patrick Kidger e2f004afd4 Added support for jax.typing.ArrayLike; now works with PyTorch's bool 2023-02-25 17:01:27 -08:00
Patrick Kidger 01f8f20bf5 Fixed a few names 2022-09-07 12:20:42 -07:00
Patrick Kidger a53fe6af57 Changes from feedback:
- Float32 introduced to replace f32 (etc.)
- Old f32 aliases were previously around for backward-compatibility, but
  the code is pretty hideous, and all stakeholders are now on board
  through RFC #13.
- Overall feedback was that Int{Sign,Unsign} was too long and not enough
  like numpy. Changed to Int, UInt, and Integer.
2022-09-07 11:09:48 -07:00
Patrick Kidger 48e8131247 Added jaxtyping.Array=jnp.ndarray 2022-08-30 12:56:47 -07:00
Patrick Kidger 903000f3d5 Rewrote syntax 2022-08-30 12:47:45 -07:00
Patrick Kidger 07259aa5c8 Added support for symbolic expressions in axes. Closes #9. 2022-08-01 01:00:18 +01:00
Patrick Kidger 6ff4620d1a Broadcast+variadic is now more lenient, and supports anything broadcastable together. Closes #8 2022-07-31 23:59:20 +01:00
Patrick Kidger 6e2837e5b7 Tidied up _#* into acting as prefixes, regardless of order. Closes #10 2022-07-31 23:45:40 +01:00
Patrick Kidger 9f75958b2d Make test/helpers.py a relative import 2022-07-12 10:03:19 +01:00
Patrick Kidger 0800e5fa38 renamed tests->test so release workflow works 2022-07-11 13:06:54 +01:00