Owen L

Sharding integration is much slower than pmap

22

This is a follow up to https://github.com/patrick-kidger/diffrax/issues/403, where when I do the equinox sharding design, I see it is substantially slower than a pmap (or even just a vmap on...

Add KL divergence terms for Latent SDEs

9

Addresses https://github.com/patrick-kidger/diffrax/issues/401. Revives https://github.com/patrick-kidger/diffrax/pull/104. Based on that PR, I made the minimal requirements to get it up to current version (e.g. taking callables instead of ODE terms since we can't...

Best way to approach KL divergence

2

I was wondering what the best way of recreating the KL divergence from `torchsde` would be in diffrax. I have a hacked together version of https://github.com/google-research/torchsde/blob/master/examples/latent_sde.py, but it seems like...

Conv Transpose doesn't exactly transpose

3

When checking if ConvTranspose is actually computing the transpose operation, it seems to be failing. I've tried different weight matrix shapes, but I'm uncertain as to why this is failing:...

Add a [source] to docs [feature request]

3

For users that are interested/regularly investigate the source code, having a method to go from docs to associated source code is very helpful (this is what I am talking about:...

### Description If I were to run a constant size function and scale it linearly across more and more CPUs/GPUs (just running the same function and the same input again...

Interest in new solvers

2

I have two classes of solvers I'm implementing with diffrax which I thought could be good to actually just put in diffrax, but I wanted to check if there was...

How to solve SDEs with NeuralPDE?

2

In the first line of the docs, it says "Physics-Informed Neural Networks for ODE, SDE, RODE, and PDE solving." but I wasn't able to find documentation/examples for SDEs like there...

Document Named Collectives

2

Named collectives have come up before in documentation for pmap/shmap, but usually links to https://jax.readthedocs.io/en/latest/jax.lax.html#parallel-operators (other relevant issue: https://github.com/google/jax/issues/13794). However, these also work for things such as vmap and I...

Spectral Norm Features

1

The original spectral norm paper didn't accurately estimate spectral norms (e.g. in the convolution case they just reshape it then do normal matrix, but this isn't the actual transpose operation,...