search

google / edward2

A simple probabilistic programming language.
Apache License 2.0
679 stars 75 forks source link

Updates MIMO's DenseMultiHead layer to support TPUs. #510

Closed copybara-service[bot] closed 2 years ago

copybara-service[bot] commented 2 years ago

Updates MIMO's DenseMultiHead layer to support TPUs.

Previously, runs would fail on TPUs with the error message Dynamic input dimension to reshape that is both splitted and combined is not supported.