17 projects
jaxtyping
Type annotations and runtime checking for shape and dtype of JAX/NumPy/PyTorch/etc. arrays.
lineax
Linear solvers in JAX and Equinox.
optimistix
Nonlinear optimisation in JAX and Equinox.
equinox
Elegant easy-to-use neural networks in JAX.
quax
Multiple dispatch in JAX via custom interpreters.
torchtyping
Runtime type annotations for the shape, dtype etc. of PyTorch Tensors.
diffrax
GPU+autodiff-capable ODE/SDE/CDE solvers written in JAX.
typst-pyimage
Typst extension, adding support for generating figures using inline Python code
pytkdocs-tweaks
Some custom tweaks to the results produced by pytkdocs (part of mkdocstrings).
sympy2jax
Turn SymPy expressions into trainable JAX expressions.
torchsde
SDE solvers and stochastic adjoint sensitivity analysis in PyTorch.
sympytorch
Turning SymPy expressions into PyTorch modules.
mkposters
Make posters from Markdown files.
mkdocs-include-exclude-files
Modify which files MkDocs includes or excludes.
torchcde
Differentiable controlled differential equation solvers for PyTorch with GPU support and memory-efficient adjoint backpropagation.
signatory
Differentiable computations of the signature and logsignature transforms, on both CPU and GPU.
py2annotate
An extension to Sphinx :code:`autodoc` to augment Sphinx documentation with type annotations, when using Python 2 style type annotations.
