Scientific computing in JAX

• sympy2jax

  1 305 4.4 Python

  Turn SymPy expressions into trainable JAX expressions.

  

sympy2jax: sympy->JAX conversion;

• jaxtyping

  7 941 8.3 Python

  Type annotations and runtime checking for shape and dtype of JAX/NumPy/PyTorch/etc. arrays. https://docs.kidger.site/jaxtyping/

  

jaxtyping: rich shape & dtype annotations for arrays and tensors (also supports PyTorch/TensorFlow/NumPy);

• InfluxDB

  www.influxdata.com sponsored

  Power Real-Time Data Analytics at Scale. Get real-time insights from all types of time series data with InfluxDB. Ingest, query, and analyze billions of data points in real-time with unbounded cardinality.

  

• eqxvision

  3 95 2.3 Python

  A Python package of computer vision models for the Equinox ecosystem.

Eqxvision: computer vision.

• diffrax

  21 1,230 8.3 Python

  Numerical differential equation solvers in JAX. Autodifferentiable and GPU-capable. https://docs.kidger.site/diffrax/

Sure. So I've got some PyTorch benchmarks here. The main take-away so far has been that for a neural ODE, the backward pass takes about 50% longer in PyTorch, and the forward (inference) pass takes an incredible 100x longer.

Suggest a related project