search

google-research / prompt-tuning

Original Implementation of Prompt Tuning from Lester, et al, 2021
Apache License 2.0
645 stars 58 forks source link

[LSC] change uses of jax.random.KeyArray and jax.random.PRNGKeyArray to jax.Array #290

Closed copybara-service[bot] closed 11 months ago

copybara-service[bot] commented 11 months ago

[LSC] change uses of jax.random.KeyArray and jax.random.PRNGKeyArray to jax.Array

This change replaces uses of jax.random.KeyArray and jax.random.PRNGKeyArray in the context of type annotations with jax.Array, which is the correct annotation for JAX PRNG keys moving forward.

The purpose of this change is to remove references to KeyArray and PRNGKeyArray, which are deprecated (https://github.com/google/jax/pull/17594) and will soon be removed from JAX. The design and thought process behind this is described in https://jax.readthedocs.io/en/latest/jep/9263-typed-keys.html.

Note that KeyArray and PRNGKeyArray have always been aliased to Any, so the new type annotation is far more specific than the old one.

github-actions[bot] commented 11 months ago

Unfortunately, we cannot accept contributions to the Prompt Tuning repo at this time. Please file issues as needed though!