rixty/KoboldAI-Client
1
1
Fork 0
mirror of https://github.com/KoboldAI/KoboldAI-Client.git synced 2025-01-24 22:19:10 +01:00
Code Issues Projects Releases Wiki Activity

Fix np.take (https://github.com/google/jax/issues/3774)

Browse Source
This commit is contained in:
Gnome Ann 2022-01-17 13:54:02 -05:00
parent 6502af086f
commit 31735c4239
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
tpu_mtj_backend.py
View File

@ -74,6 +74,7 @@ def apply_repetition_penalty_dynamic(logits, tokens, repetition_penalty):
This gets called by generate_loop_fn to apply repetition penalty
to the 1D array logits using the provided 1D array of tokens to penalize
'''
tokens = np.minimum(tokens, params["n_vocab"]-1) # https://github.com/google/jax/issues/3774
# Make a new array with the same length as the tokens array but with
# each element replaced by the value at the corresponding index in the
# logits array; e.g.