mirror of
https://github.com/wassname/jaxtyping.git
synced 2026-09-09 11:24:55 +08:00
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"].
This commit is contained in:
Patrick Kidger
@@ -0,0 +1,12 @@
|
||||
If you found this library to be useful in academic work, then please cite: ([arXiv link](https://arxiv.org/abs/2111.00254))
|
||||
|
||||
```bibtex
|
||||
@article{kidger2021equinox,
|
||||
author={Patrick Kidger and Cristian Garcia},
|
||||
title={{E}quinox: neural networks in {JAX} via callable {P}y{T}rees and filtered transformations},
|
||||
year={2021},
|
||||
journal={Differentiable Programming workshop at Neural Information Processing Systems 2021}
|
||||
}
|
||||
```
|
||||
|
||||
(Also consider starring the project [on GitHub](https://github.com/patrick-kidger/equinox).)
|
||||
@@ -0,0 +1 @@
|
||||
ErrorDocument 404 /jaxtyping/404.html
|
||||
@@ -0,0 +1,20 @@
|
||||
{% import "partials/language.html" as lang with context %}
|
||||
<a href="{{ config.repo_url }}" title="{{ lang.t('source.link.title') }}" class="md-source" data-md-component="source">
|
||||
<div class="md-source__icon md-icon">
|
||||
{% set icon = config.theme.icon.repo or "fontawesome/brands/git-alt" %}
|
||||
{% include ".icons/" ~ icon ~ ".svg" %}
|
||||
</div>
|
||||
<div class="md-source__repository">
|
||||
{{ config.repo_name }}
|
||||
</div>
|
||||
</a>
|
||||
{% if config.theme.twitter_url %}
|
||||
<a href="{{ config.theme.twitter_url }}" title="Go to Twitter" class="md-source">
|
||||
<div class="md-source__icon md-icon">
|
||||
{% include ".icons/fontawesome/brands/twitter.svg" %}
|
||||
</div>
|
||||
<div class="md-source__repository">
|
||||
{{ config.theme.twitter_name }}
|
||||
</div>
|
||||
</a>
|
||||
{% endif %}
|
||||
Vendored
+2
@@ -0,0 +1,2 @@
|
||||
The favicon is `math-integral` from https://materialdesignicons.com, found by way of https://pictogrammers.com.
|
||||
(The logo is `math-integral-box`.)
|
||||
Vendored
+167
@@ -0,0 +1,167 @@
|
||||
/* Fix /page#foo going to the top of the viewport and being hidden by the navbar */
|
||||
html {
|
||||
scroll-padding-top: 50px;
|
||||
}
|
||||
|
||||
/* Fit the Twitter handle alongside the GitHub one in the top right. */
|
||||
|
||||
div.md-header__source {
|
||||
width: revert;
|
||||
max-width: revert;
|
||||
}
|
||||
|
||||
a.md-source {
|
||||
display: inline-block;
|
||||
}
|
||||
|
||||
.md-source__repository {
|
||||
max-width: 100%;
|
||||
}
|
||||
|
||||
/* Emphasise sections of nav on left hand side */
|
||||
|
||||
nav.md-nav {
|
||||
padding-left: 5px;
|
||||
}
|
||||
|
||||
nav.md-nav--secondary {
|
||||
border-left: revert !important;
|
||||
}
|
||||
|
||||
.md-nav__title {
|
||||
font-size: 0.9rem;
|
||||
}
|
||||
|
||||
.md-nav__item--section > .md-nav__link {
|
||||
font-size: 0.9rem;
|
||||
}
|
||||
|
||||
/* Indent autogenerated documentation */
|
||||
|
||||
div.doc-contents {
|
||||
padding-left: 25px;
|
||||
border-left: 4px solid rgba(230, 230, 230);
|
||||
}
|
||||
|
||||
/* Increase visibility of splitters "---" */
|
||||
|
||||
[data-md-color-scheme="default"] .md-typeset hr {
|
||||
border-bottom-color: rgb(0, 0, 0);
|
||||
border-bottom-width: 1pt;
|
||||
}
|
||||
|
||||
[data-md-color-scheme="slate"] .md-typeset hr {
|
||||
border-bottom-color: rgb(230, 230, 230);
|
||||
}
|
||||
|
||||
/* More space at the bottom of the page */
|
||||
|
||||
.md-main__inner {
|
||||
margin-bottom: 1.5rem;
|
||||
}
|
||||
|
||||
/* Remove prev/next footer buttons */
|
||||
|
||||
.md-footer__inner {
|
||||
display: none;
|
||||
}
|
||||
|
||||
/* Change font sizes */
|
||||
|
||||
html {
|
||||
/* Decrease font size for overall webpage
|
||||
Down from 137.5% which is the Material default */
|
||||
font-size: 110%;
|
||||
}
|
||||
|
||||
.md-typeset .admonition {
|
||||
/* Increase font size in admonitions */
|
||||
font-size: 100% !important;
|
||||
}
|
||||
|
||||
.md-typeset details {
|
||||
/* Increase font size in details */
|
||||
font-size: 100% !important;
|
||||
}
|
||||
|
||||
.md-typeset h1 {
|
||||
font-size: 1.6rem;
|
||||
}
|
||||
|
||||
.md-typeset h2 {
|
||||
font-size: 1.5rem;
|
||||
}
|
||||
|
||||
.md-typeset h3 {
|
||||
font-size: 1.3rem;
|
||||
}
|
||||
|
||||
.md-typeset h4 {
|
||||
font-size: 1.1rem;
|
||||
}
|
||||
|
||||
.md-typeset h5 {
|
||||
font-size: 0.9rem;
|
||||
}
|
||||
|
||||
.md-typeset h6 {
|
||||
font-size: 0.8rem;
|
||||
}
|
||||
|
||||
/* Bugfix: remove the superfluous parts generated when doing:
|
||||
|
||||
??? Blah
|
||||
|
||||
::: library.something
|
||||
*/
|
||||
|
||||
.md-typeset details .mkdocstrings > h4 {
|
||||
display: none;
|
||||
}
|
||||
|
||||
.md-typeset details .mkdocstrings > h5 {
|
||||
display: none;
|
||||
}
|
||||
|
||||
/* Change default colours for <a> tags */
|
||||
|
||||
[data-md-color-scheme="default"] {
|
||||
--md-typeset-a-color: rgb(0, 189, 164) !important;
|
||||
}
|
||||
[data-md-color-scheme="slate"] {
|
||||
--md-typeset-a-color: rgb(0, 189, 164) !important;
|
||||
}
|
||||
|
||||
/* Highlight functions, classes etc. type signatures. Really helps to make clear where
|
||||
one item ends and another begins. */
|
||||
|
||||
[data-md-color-scheme="default"] {
|
||||
--doc-heading-color: #DDD;
|
||||
--doc-heading-border-color: #CCC;
|
||||
--doc-heading-color-alt: #F0F0F0;
|
||||
}
|
||||
[data-md-color-scheme="slate"] {
|
||||
--doc-heading-color: rgb(25,25,33);
|
||||
--doc-heading-border-color: rgb(25,25,33);
|
||||
--doc-heading-color-alt: rgb(33,33,44);
|
||||
--md-code-bg-color: rgb(38,38,50);
|
||||
}
|
||||
|
||||
h4.doc-heading {
|
||||
/* NOT var(--md-code-bg-color) as that's not visually distinct from other code blocks.*/
|
||||
background-color: var(--doc-heading-color);
|
||||
border: solid var(--doc-heading-border-color);
|
||||
border-width: 1.5pt;
|
||||
border-radius: 2pt;
|
||||
padding: 0pt 5pt 2pt 5pt;
|
||||
}
|
||||
h5.doc-heading, h6.heading {
|
||||
background-color: var(--doc-heading-color-alt);
|
||||
border-radius: 2pt;
|
||||
padding: 0pt 5pt 2pt 5pt;
|
||||
}
|
||||
|
||||
/* Make errors in notebooks have scrolling */
|
||||
.output_error > pre {
|
||||
overflow: auto;
|
||||
}
|
||||
BIN
Binary file not shown.
|
After Width: | Height: | Size: 807 B |
Vendored
+16
@@ -0,0 +1,16 @@
|
||||
window.MathJax = {
|
||||
tex: {
|
||||
inlineMath: [["\\(", "\\)"]],
|
||||
displayMath: [["\\[", "\\]"]],
|
||||
processEscapes: true,
|
||||
processEnvironments: true
|
||||
},
|
||||
options: {
|
||||
ignoreHtmlClass: ".*|",
|
||||
processHtmlClass: "arithmatex"
|
||||
}
|
||||
};
|
||||
|
||||
document$.subscribe(() => {
|
||||
MathJax.typesetPromise()
|
||||
})
|
||||
@@ -0,0 +1,13 @@
|
||||
# Advanced features
|
||||
|
||||
## Abstract base classes
|
||||
|
||||
::: jaxtyping.AbstractDtype
|
||||
selection:
|
||||
members:
|
||||
false
|
||||
|
||||
::: jaxtyping.AbstractArray
|
||||
selection:
|
||||
members:
|
||||
false
|
||||
@@ -0,0 +1,112 @@
|
||||
# Array annotations
|
||||
|
||||
The shape and dtypes of arrays can be annotated in the form `dtype[array, shape]`, such as `Float[Array, "batch channels"]`.
|
||||
|
||||
## Shape
|
||||
|
||||
**Symbols**
|
||||
|
||||
The shape should be a string of space-separated symbols, such as `"a b c d"`. Each symbol can be either an:
|
||||
|
||||
- `int`: fixed-size axis, e.g. `"28 28"`.
|
||||
- `str`: variable-size axis, e.g. `"channels"`.
|
||||
- A symbolic expression (without spaces!) in terms of other variable-size axes, e.g.
|
||||
`def remove_last(x: Float[Array, "dim"]) -> Float[Array, "dim-1"]`.
|
||||
|
||||
When calling a function, variable-size axes and symbolic axes will be matched up across all arguments and checked for consistency. (See [Runtime type checking](./runtime-type-checking.md).)
|
||||
|
||||
**Modifiers**
|
||||
|
||||
In addition some modifiers can be applied:
|
||||
|
||||
- Prepend `*` to a dimension to indicate that it can match multiple axes, e.g. `"*batch c h w"` will match zero or more batch axes.
|
||||
- Prepend `#` to a dimension to indicate that it can be that size *or* equal to one -- i.e. broadcasting is acceptable, e.g.
|
||||
`def add(x: Float[Array, "#foo"], y: Float[Array, "#foo"]) -> Float[Array, "#foo"]`.
|
||||
- Prepend `_` to a dimension to disable any runtime checking of that dimension (so that it can be used just as documentation). This can also be used as just `_` on its own: e.g. `"b c _ _"`.
|
||||
- Documentation-only names (i.e. they're ignored by jaxtyping) can be handled by prepending a name followed by `=` e.g. `Float[Array, "rows=4 cols=3"]`.
|
||||
|
||||
When using multiple modifiers, their order does not matter.
|
||||
|
||||
As a special case:
|
||||
|
||||
- `...`: anonymous zero or more axes (equivalent to `*_`) e.g. `"... c h w"`
|
||||
|
||||
**Notes**
|
||||
|
||||
- To denote a scalar shape use `""`, e.g. `Float[Array, ""]`.
|
||||
- To denote an arbitrary shape (and only check dtype) use `"..."`, e.g. `Float[Array, "..."]`.
|
||||
- You cannot have more than one use of multiple-axes, i.e. you can only use `...` or `*name` at most once in each array.
|
||||
- An example of broadcasting multiple dimensions:
|
||||
`def add(x: Float[Array, "*#foo"], y: Float[Array, "*#foo"]) -> Float[Array, "*#foo"]`.
|
||||
- A symbolic expression cannot be evaluated unless all of the axes sizes it refers to have already been processed. In practice this usually means that they should only be used in annotations for the return type, and only use axes declared in the arguments.
|
||||
|
||||
## Dtype
|
||||
|
||||
The dtype should be any one of (all imported from `jaxtyping`):
|
||||
|
||||
- Any dtype at all: `Shaped`
|
||||
- Boolean: `Bool`
|
||||
- PRNG key: `Key`
|
||||
- Any integer, unsigned integer, floating, or complex: `Num`
|
||||
- Any floating or complex: `Inexact`
|
||||
- Any floating point: `Float`
|
||||
- Of particular precision: `BFloat16`, `Float16`, `Float32`, `Float64`
|
||||
- Any complex: `Complex`
|
||||
- Of particular precision: `Complex64`, `Complex128`
|
||||
- Any integer or unsigned intger: `Integer`
|
||||
- Any unsigned integer: `UInt`
|
||||
- Of particular precision: `UInt8`, `UInt16`, `UInt32`, `UInt64`
|
||||
- Any signed integer: `Int`
|
||||
- Of particular precision: `Int8`, `Int16`, `Int32`, `Int64`
|
||||
|
||||
Unless you really want to force a particular precision, then for most applications you should probably allow any floating-point, any integer, etc. That is, use
|
||||
```python
|
||||
from jaxtyping import Array, Float
|
||||
Float[Array, "some_shape"]
|
||||
```
|
||||
rather than
|
||||
```python
|
||||
from jaxtyping import Array, Float32
|
||||
Float32[Array, "some_shape"]
|
||||
```
|
||||
|
||||
## Array
|
||||
|
||||
The array should usually be a `jaxtyping.Array`, which is an alias for `jax.numpy.ndarray` (which is itself an alias for `jax.Array`).
|
||||
|
||||
`jaxtyping.ArrayLike` is also available, which is an alias for `jax.typing.ArrayLike`. This is a union over JAX arrays and the builtin `bool`/`int`/`float`/`complex`.
|
||||
|
||||
You can use non-JAX types as well. jaxtyping also supports NumPy, TensorFlow, and PyTorch, e.g.:
|
||||
```python
|
||||
Float[np.ndarray, "..."]
|
||||
Float[tf.Tensor, "..."]
|
||||
Float[torch.Tensor, "..."]
|
||||
```
|
||||
|
||||
Shape-and-dtype specified jaxtyping arrays can also be used, e.g.
|
||||
```python
|
||||
Image = Float[Array, "channels height width"]
|
||||
BatchImage = Float[Image, "batch"]
|
||||
```
|
||||
in which case the additional shape is prepended, and the acceptable dtypes are the intersection of the two dtype specifiers used. (So that e.g. `BatchImage = Shaped[Image, "batch"]` would work just as well. But `Bool[Image, "batch"]` would throw an error, as there are no dtypes that are both bools and floats.) Thus the above is equivalent to
|
||||
```python
|
||||
BatchImage = Float[Array, "batch channels height width"]
|
||||
```
|
||||
|
||||
Note that `jaxtyping.{Array, ArrayLike}` are only available if JAX has been installed.
|
||||
|
||||
## Scalars, PRNGKeys
|
||||
|
||||
For convenience, jaxtyping also includes `jaxtyping.Scalar`, `jaxtyping.ScalarLike`, and `jaxtyping.PRNGKeyArray`, defined as:
|
||||
```python
|
||||
Scalar = Shaped[Array, ""]
|
||||
ScalarLike = Shaped[ArrayLike, ""]
|
||||
|
||||
# Depending on the value of `JAX_ENABLE_CUSTOM_PRNG`:
|
||||
PRNGKeyArray = Key[Array, ""]
|
||||
PRNGKeyArray = UInt32[Array, "2"]
|
||||
```
|
||||
|
||||
Recalling that shape-and-dtype specified jaxtyping arrays can be nested, this means that e.g. you can annotate the output of `jax.random.split` with `Shaped[PRNGKeyArray, "2"]`, or e.g. an integer scalar with `Int[Scalar, ""]`.
|
||||
|
||||
Note that `jaxtyping.{Scalar, ScalarLike, PRNGKeyArray}` are only available if JAX has been installed.
|
||||
@@ -0,0 +1,8 @@
|
||||
# PyTree annotations
|
||||
|
||||
:::jaxtyping.PyTree
|
||||
selection:
|
||||
members:
|
||||
false
|
||||
|
||||
Note that `jaxtyping.PyTree` is only available if JAX has been installed.
|
||||
@@ -0,0 +1,22 @@
|
||||
# Runtime type checking
|
||||
|
||||
(See the [FAQ](../faq.md) for details on static type checking.)
|
||||
|
||||
Runtime type checking **synergises beautifully with `jax.jit`!** All shape checks will be performed at trace-time only, and will not impact runtime performance.
|
||||
|
||||
Runtime type-checking should be performed using a library like [typeguard](https://github.com/agronholm/typeguard) or [beartype](https://github.com/beartype/beartype).
|
||||
|
||||
The types provided by `jaxtyping`, e.g. `Float[Array, "batch channels"]`, are all compatible with `isinstance` checks, e.g. `isinstance(x, Float[Array, "batch channels"])`. This means that jaxtyping should be compatible with all runtime type checkers out-of-the-box.
|
||||
|
||||
Some additional context is needed to ensure consistency between multiple argments (i.e. that shapes match up between arrays). For this, you can use either `jaxtyping.jaxtyped` to add this capability to a single function, or `jaxtyping.install_import_hook` to add this capability to a whole codebase. If either are too much magic for you, you can safely use neither and have just single-argument type checking.
|
||||
|
||||
::: jaxtyping.jaxtyped
|
||||
|
||||
---
|
||||
|
||||
It can be a lot of effort to add `@jaxtyped` decorators all over your codebase.
|
||||
(Not to mention that double-decorators everywhere are a bit ugly.)
|
||||
|
||||
The easier option is usually to use the import hook.
|
||||
|
||||
::: jaxtyping.install_import_hook
|
||||
+48
@@ -0,0 +1,48 @@
|
||||
# FAQ
|
||||
|
||||
## Is jaxtyping compatible with static type checkers like `mypy`/`pyright`/`pytype`?
|
||||
|
||||
There is partial support for these. An annotation of the form `dtype[array, shape]` should be treated as just `array` by a static type checker. Unfortunately full dtype/shape checking is beyond the scope of what static type checking is currently capable of.
|
||||
|
||||
(Note that at time of writing, `pytype` has a bug in that `dtype[array, shape]` is sometimes treated as `Any` rather than `array`. `mypy` and `pyright` both work fine.)
|
||||
|
||||
## How does jaxtyping interact with `jax.jit`?
|
||||
|
||||
jaxtyping and `jax.jit` synergise beautifully.
|
||||
|
||||
When calling JAX operations wrapped in a `jax.jit`, then the dtype/shape-checking will happen at trace time. (When JAX traces your function prior to compiling it.) The actual compiled code does not have any dtype/shape-checking, and will therefore still be just as fast as before!
|
||||
|
||||
## `flake8` or Ruff are throwing an error.
|
||||
|
||||
In type annotations, strings are used for two different things. Sometimes they're strings. Sometimes they're "forward references", used to refer to a type that will be defined later.
|
||||
|
||||
Some tooling in the Python ecosystem assumes that only the latter is true, and will throw spurious errors if you try to use a string just as a string (like we do).
|
||||
|
||||
In the case of `flake8`, or Ruff, this can be resolved. Multi-dimensional arrays (e.g. `Float32[Array, "b c"]`) will throw a very unusual error (F722, syntax error in forward annotation), so you can safely just disable this particular error globally. Uni-dimensional arrays (e.g. `Float32[Array, "x"]`) will throw an error that's actually useful (F821, undefined name), so instead of disabling this globally, you should instead prepend a space to the start of your shape, e.g. `Float32[Array, " x"]`. `jaxtyping` will treat this in the same way, whilst `flake8` will now throw an F722 error that you can disable as before.
|
||||
|
||||
## Does jaxtyping use [PEP 646](https://www.python.org/dev/peps/pep-0646/) (variadic generics)?
|
||||
|
||||
The intention of PEP 646 was to make it possible for static type checkers to perform shape checks of arrays. Unfortunately, this still isn't yet practical, so jaxtyping deliberately does not use this. (Yet?)
|
||||
|
||||
The real problem is that Python's static typing ecosystem is a complicated collection of edge cases. Many of them block ML/scientific computing in particular. For example:
|
||||
|
||||
1. The static type system is intrinsically not expressive enough to describe operations like concatenation, stacking, or broadcasting.
|
||||
|
||||
2. Axes have to be lifted to type-level variables. Meanwhile the approach taken in libraries like `jaxtyping` and [TorchTyping](https://github.com/patrick-kidger/torchtyping) is to use value-level variables for types: because that's what the underlying JAX, PyTorch etc. libraries use! As such, making a static type checker work with these libraries would require either fundamentally rewriting these libraries, or exhaustively maintaining type stubs for them, and would *still* require a `typing.cast` any time you use anything unstubbed (e.g. any third party library, or part of your codebase you haven't typed yet). This is a huge maintenance burden.
|
||||
|
||||
3. Static type checkers have a variety of bugs that affect this use case. `mypy` doesn't support `Protocol`s correctly. `pyright` doesn't support genericised subprotocols. etc.
|
||||
|
||||
4. Variadic generics exist. Variadic protocols do not. (It's not clear that these were contemplated.)
|
||||
|
||||
5. The syntax for static typing is a little verbose. You have to write things like `Array[Float32, Unpack[AnyShape], Literal[3], Height, Width]` instead of `Float32[Array, "... 3 height width"]`.
|
||||
|
||||
6. [The underlying type system has flaws](https://github.com/patrick-kidger/torchtyping/issues/37#issuecomment-1153294196).
|
||||
[The numeric tower is broken](https://stackoverflow.com/a/69383462);
|
||||
[int is not a number](https://github.com/python/mypy/issues/3186#issuecomment-885718629);
|
||||
[virtual base classes don't work](https://github.com/python/mypy/issues/2922);
|
||||
[complex lies about having comparison operations, so type checkers have to lie about that lie in order to remove them again](https://posita.github.io/numerary/0.4/whytho/);
|
||||
`typing.*` don't work with `isinstance`;
|
||||
co/contra-variance are baked into containers (not specified at use-time);
|
||||
`dict` is variadic despite... not being variadic;
|
||||
bool is a subclass of int (!);
|
||||
... etc. etc.
|
||||
@@ -0,0 +1,42 @@
|
||||
# Getting started
|
||||
|
||||
jaxtyping is a library providing type annotations **and runtime type-checking** for:
|
||||
|
||||
1. shape and dtype of [JAX](https://github.com/google/jax) arrays;
|
||||
2. [PyTrees](https://jax.readthedocs.io/en/latest/pytrees.html).
|
||||
|
||||
*(Now also supports PyTorch, NumPy, and TensorFlow!)*
|
||||
|
||||
## Installation
|
||||
|
||||
```bash
|
||||
pip install jaxtyping
|
||||
```
|
||||
|
||||
Requires Python 3.8+.
|
||||
|
||||
JAX is an optional dependency, required for a few JAX-specific types. If JAX is not installed then these will not be available, but you may still use jaxtyping to provide shape/dtype annotations for PyTorch/NumPy/TensorFlow/etc.
|
||||
|
||||
The annotations provided by jaxtyping are compatible with runtime type-checking packages, so it is common to also install one of these. The two most popular are [typeguard](https://github.com/agronholm/typeguard) (which exhaustively checks every argument) and [beartype](https://github.com/beartype/beartype) (which checks random pieces of arguments).
|
||||
|
||||
## Example
|
||||
|
||||
```python
|
||||
from jaxtyping import Array, Float, PyTree
|
||||
|
||||
# Accepts floating-point 2D arrays with matching dimensions
|
||||
def matrix_multiply(x: Float[Array, "dim1 dim2"],
|
||||
y: Float[Array, "dim2 dim3"]
|
||||
) -> Float[Array, "dim1 dim3"]:
|
||||
...
|
||||
|
||||
def accepts_pytree_of_ints(x: PyTree[int]):
|
||||
...
|
||||
|
||||
def accepts_pytree_of_arrays(x: PyTree[Float[Array, "batch c1 c2"]]):
|
||||
...
|
||||
```
|
||||
|
||||
## Next steps
|
||||
|
||||
Have a read of the [Array annotations](./api/array.md) documentation on the left-hand bar!
|
||||
@@ -0,0 +1,13 @@
|
||||
mkdocs==1.3.0 # Main documentation generator.
|
||||
mkdocs-material==7.3.6 # Theme
|
||||
pymdown-extensions==9.4 # Markdown extensions e.g. to handle LaTeX.
|
||||
mkdocstrings==0.17.0 # Autogenerate documentation from docstrings.
|
||||
mknotebooks==0.7.1 # Turn Jupyter Lab notebooks into webpages.
|
||||
pytkdocs_tweaks==0.0.5 # Tweaks mkdocstrings to improve various aspects
|
||||
mkdocs_include_exclude_files==0.0.1 # Tweak which files are included/excluded
|
||||
jinja2==3.0.3 # Older version. After 3.1.0 seems to be incompatible with current versions of mkdocstrings.
|
||||
pygments==2.14.0
|
||||
|
||||
# Dependencies of jaxtyping itself.
|
||||
# Always use most up-to-date versions.
|
||||
jax[cpu]
|
||||
Reference in New Issue
Block a user