Wilson Yan

icon indicating a website link https://wilson1yan.github.io/ icon indicating an email [email protected]

icon company University of California, Berkeley icon location Berkeley, CA icon location PhD student interested in generative modeling and representation learning

Results 3 issues of Wilson Yan

Different runtimes for `pmap` vs `pjit`

4 comment

Hi, I have a simple script below that compares runtimes for `pmap` vs `pjit`. I expected that the runtime for `pjit` with full data parallelism would be the same for...

bug

Possible bug in StyleGAN training code

1 comment

Hi, Nice repo! Is this line a bug? Since I think `batch['images']` is `N x B x H x W x C`, so the indices should be shift up by...

jnp.cumsum is incorrect when sharding over summing axis on GPUs

### Description I have this simple script below that tests `jnp.cumsum` when sharding along the same axis it is summing over. Tested on a machine with 8 40GB A100s. ```python...

bug