flax
527cb4a4 - Merge pull request #4649 from IvyZX:attn

Commit

1 year ago

Merge pull request #4649 from IvyZX:attn PiperOrigin-RevId: 740394095

References

#4649 - Use jax.nn.dot_product_attention when possible

Author

a-googler

a-googler

Parents

Loading