Bill Mark

Add documentation on how to use PyType with Jax (and also common add-on libraries such as Flax)

10

There are a bunch of tricks that one needs to know to use PyType with JAX (esp. also in combination with Flax). For example, a PyTree needs to be treated...

Improve RTD for initializers

7

Provide as much information as possible. At least, this should include a description of your issue and steps to reproduce the problem. If possible also provide a summary of what...

Checkify exception only reports error from one device (not all)

7

### Description When running with four hosts and four devices on each host, I see an "errs" returned by pmap of checkify that looks like the folllowing: ``` Error(at mapped...

jax.debug.breakpoint gives UnexpectedTracerError when used with jax.lax.cond

3

### Description ``` import jax import jax.numpy as jnp def f(x, example): jax.lax.cond(example == 1, jax.debug.breakpoint, lambda *args: None) return x f_vmap = jax.vmap(f, in_axes=(0, None), out_axes=0) def g(x, example):...