Paul Balanca

WIP: Implement a `denormalize` custom Jaxpr operator simplifying MCX logpdfs

2

# Overview This PR is implementing a generic `denormalize` decorator which removes normalizing constants in a logpdf. Per call to contributions in #65. # Implementation The current implementation is a...

It would be very useful to be able to use directly `batch_bool` mask to increment/decrement integer `batch`, using the integer overfull trick. It allows to update a table of indices...

L2 weight decay implementation in TF benchmarks

4

The existing implementation use a L2 regularization on all model variables (including batch norm variables and biases). It's quite different from TF slim models which usually regularizes only conv2d weights,...

`jnp.frexp` not matching Numpy on sub-normals

2

### Description The function `jnp.frexp` does not match Numpy result on sub-normals (I would personally consider the Numpy result to be the more accurate one). ```python import numpy as np...

### Description The `jax.debug.callback` IPU unit tests is blocking when ignoring the returned value of the jitted function. Probably a deadlock in the IPU PjRt client at shutdown (buffer deleted?)....

`import jax` is crashing with a cryptic message when `jaxlib` version does not coincide with the local Poplar SD (see #12 ) We should implement a mechanism for early catching...