Giles Harper-Donnelly

Error when inheriting from JointDistributionSequential in tfp version 0.18.0

3

In the most recent update (0.18.0) instantiating a class which inherits from `tfp.distributions.JointDistributionSequential` now raises an error. Here is a MWE (reproduced in [colab](https://colab.research.google.com/drive/1gG_I3gDdbMklfEobqWMuCFa9tv-SR9Dj?usp=sharing)): ```python from jax import numpy as...

http://forneylab.org/ domain expired

1

The [ForneyLab project page](http://forneylab.org/) linked in the README links to an expired domain. Thought you might want a heads up :)

Closes #12377. As discussed in #12377 unlike similar functions in `jax.tree_util` the function `tree_reduce` does not provide the option to pass in an `is_leaf` callable. This PR adds that functionality...

Add the ability to pass an `is_leaf` argument to `tree_reduce`.

2

Many other functions in `jax.tree_util` allow the user to provide an `is_leaf` argument which is subsequently passed to `tree_flatten`. It would be nice to be able to pass such a...

Error initializing tfb.AutoregressiveNetwork using jax substrate

1

Trying to initialize an instance of `tfb.AutoregressiveNetwork` using the jax substrate fails with an AttributeError. With the example usage from the [docs](https://www.tensorflow.org/probability/api_docs/python/tfp/substrates/jax/bijectors/AutoregressiveNetwork): ```python from tensorflow_probability.substrates import jax as tfp tfb...

Add JAX kmeans implementation

1

## Summary This PR add an implementation of the k-means algorithm written in JAX. By replacing the current calls to the `sklearn` implementation with this new implementation, we can now...

# What does this PR do? The [Migrating from Haiku/Linen to NNX](https://flax.readthedocs.io/en/latest/nnx/haiku_linen_vs_nnx.html) page in the docs contains a broken link to the [NNX Basics](https://flax.readthedocs.io/en/latest/nnx/nnx_basics.html) page. This PR updates the link...

## Summary This PR is a Work In Progress. Fix the current errors with type hints and some common linting errors. ## Details ### Commits A brief categorisation of the...

## Summary This PR changes the parameters of the `LinearRegressionHMM` test to prevent it failing. ## Details At present the `LinearRegressionHMM` `test_sample_and_fit` test is failing. The `monotonically_increasing` assertion is failing...