Skip to content

GRU pruning example#515

Open
SheilaSeidelADI wants to merge 12 commits intoVainF:masterfrom
SheilaSeidelADI:gru-example
Open

GRU pruning example#515
SheilaSeidelADI wants to merge 12 commits intoVainF:masterfrom
SheilaSeidelADI:gru-example

Conversation

@SheilaSeidelADI
Copy link
Copy Markdown

@SheilaSeidelADI SheilaSeidelADI commented Aug 20, 2025

Pull Request Overview

This PR introduces GRU (Gated Recurrent Unit) pruning capabilities for the torch-pruning library. The key innovation is a custom PrunableGRU implementation (composed of more friendly prunable units) that enables torch-pruning to safely graph and modify GRU dimensions.

  • Custom PrunableGRU implementation that exposes GRU operations as prunable PyTorch layers
  • Identity layer solution to break circular dependencies in recurrent hidden states
  • Conversion utilities for seamless integration between torch.nn.GRU and PrunableGRU

The idea behind this approach could be applied to other types of recursive units in the future, like LSTMs.

The relevant tickets in the repo:
#44
#404
#276
#476

Reviewed Changes

File Description
examples/GRUs/test_gru.py Unit tests validating PrunableGRU equivalence with torch.nn.GRU and testing conversion utilities
examples/GRUs/gru_utils.py Core implementation of PrunableGRU class and conversion utilities between torch.nn.GRU and PrunableGRU
examples/GRUs/gru_pruning_example.py Complete end-to-end demonstration of the GRU pruning workflow
examples/GRUs/README.md Documentation explaining the approach, limitations, and usage instructions

SheilaSeidelADI and others added 12 commits July 24, 2025 19:25
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant