17 projects
lineax
Linear solvers in JAX and Equinox.
equinox
Elegant easy-to-use neural networks in JAX.
jaxtyping
Type annotations and runtime checking for shape and dtype of JAX arrays, and PyTrees.
quax
Multiple dispatch in JAX via custom interpreters.
diffrax
GPU+autodiff-capable ODE/SDE/CDE solvers written in JAX.
optimistix
Nonlinear optimisation in JAX and Equinox.
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.
typst-pyimage
Typst extension, adding support for generating figures using inline Python code
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.
torchtyping
Runtime type annotations for the shape, dtype etc. of PyTorch Tensors.
py2annotate
An extension to Sphinx :code:`autodoc` to augment Sphinx documentation with type annotations, when using Python 2 style type annotations.
