deel-ai / deel-lip

Build and train Lipschitz constrained networks: TensorFlow implementation of k-Lipschitz layers
https://deel-ai.github.io/deel-lip/
MIT License
89 stars 10 forks source link

Feat/group sort tpu #44

Closed thib-s closed 2 years ago

thib-s commented 2 years ago

GroupSort now works on colab's TPU with a limited overhead ( 24s/epoch for GroupSort2 vs 20s/epoch for ReLU ) Also added tests for GroupSort