Support for block sparse matrices

System information

• TensorFlow version (you are using): 2.3
• Are you willing to contribute it (Yes/No): I can write tutorials.

Motivation

Sparsity is a well-studied topic in neural networks that is highly relevant from both research and engineering fronts. Model optimization techniques like pruning are fueled by the idea of sparsity. Moreover, it plays even better when combined with techniques like quantization, and knowledge distillation.

tfmot provides us with good support for pruning by allowing it from many different flavors like training models with pruning schedules from scratch instead of re-training them, ability to customize pruning mechanics within a model at a layer-level, etc.

OpenAI folks propose block sparse kernels as an improvement/alternative to rigorous sparse kernels and they show that it can greatly improve the computation time. Here's the original blog post. Here's the original paper.

As examples:

1. Instead of using Dense layers use block sparse layers in the first place. Here's some motivation from Hugging Face.