JTan

Low MIG and AAM metrics

6

Hello, Firstly, just wanted to state that this is a great repo with a very understandable code base! I seem to be getting extremely low MIG / AAM scores (around...

I would like to start implementing selective compression so that the network can be instructed to compress/synthesize certain regions while preserving selected regions. (Section 4.2 of https://arxiv.org/pdf/1804.02958.pdf). This will require...

Nesting `vmap` within `pmap` using `jax.pure_callback` segfaults

1

### Description Some background, I'm trying to parallelize a CPU-intensive computation using a callback to some `scipy.optimize` routines using `jax.pure_callback` across the available CPUs on my machine. Using `vmap` and...

Using `vmap` for root finding with a vector of parameters

8

Hi devs, I'm trying to `vmap` over an root finding procedure involving vector-valued parameters. I gather this is currently not possible with `jaxopt.ScipyRootFinding` because this is a wrapper around `scipy`...

Including user-defined Jacobian

4

Hi devs, looks like a really nice library. I've been looking for a Jax-native root finding method that supports `vmap` for some time. Currently I am using an external call...