eadadi

vim_ai_complete_default and vim_ai_edit_default model can't be changed

10

When changing the model of either vim_ai_complete_default or vim_ai_edit_default to another one (e.g. gpt-4o) then they stop working, and an error suggesting that API key should be verified is logged.

atari100k:atari_pong fails on Batcher creation

5

Running with: ```bash python3 ./recall2imagine/train.py / --logdir /logdir/$(date +%Y%m%d-%H%M%S) / --configs atari100k tiny --task atari_pong --jax.platform=cpu ``` Fails with: ``` AttributeError: 'function' object has no attribute 'serializer' ``` Traceback: ```...

Add modify option before executing shell commands

5

# Motivation Sometimes we need to make modifications to the suggested prompt before executing it. This patch adds [M]odify option to the suggested prompt

Replace functools.partial with jax.tree_util.Partial

3

Updated various buffer files to replace functools.partial with jax.tree_util.Partial for consistency and improved functionality. The motivation is to be able to use jax transformations over buffers. For example, before this...

Bug- False "UserWarning: Explicitly requested dtype <class 'jax.numpy.int64'> requested in astype is not available, and will be truncated to dtype int32

1

```python Python 3.9.19 (main, May 6 2024, 19:43:03) [GCC 11.2.0] :: Anaconda, Inc. on linux Type "help", "copyright", "credits" or "license" for more information. >>> import jax >>> import tensorflow_probability.substrates.jax...