FlaxGPTNeo (#12493)

* flax gpt neo

* fix query scaling

* update generation test

* use flax model for test
This commit is contained in:
Suraj Patil
2021-07-06 18:55:18 +05:30
committed by GitHub
parent 626a0a0147
commit 7a259c190c
8 changed files with 1064 additions and 3 deletions

View File

@@ -71,3 +71,16 @@ GPTNeoForSequenceClassification
.. autoclass:: transformers.GPTNeoForSequenceClassification
:members: forward
FlaxGPTNeoModel
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. autoclass:: transformers.FlaxGPTNeoModel
:members: __call__
FlaxGPTNeoForCausalLM
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. autoclass:: transformers.FlaxGPTNeoForCausalLM
:members: __call__