Implement HardShrink, SoftShrink and Shrink Activations

[aditya0by0]

Pull Request Template

Checklist

• Confirmed that cargo run-checks command has been executed.
• Made sure the book is up to date with changes in this PR.

Related Issues/PRs

These activations are needed to implement Shrink ONNX operator tracel-ai/burn-onnx#158

Changes

Soft and Hard Shrink Activations

Testing

• test_soft_shrink_forward
• test_soft_shrink_with_lambd_and_bias
• display
• test_hard_shrink_forward
• test_hard_shrink_forward_with_lambd
• display

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 61.32%. Comparing base (5923b1e) to head (fb17ef9).
⚠️ Report is 4 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #4556      +/-   ##
==========================================
+ Coverage   61.21%   61.32%   +0.11%     
==========================================
  Files        1062     1066       +4     
  Lines      136012   136410     +398     
==========================================
+ Hits        83258    83655     +397     
- Misses      52754    52755       +1     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[aditya0by0]

@antimora, Please review.

[Copilot]

Pull request overview

This PR implements HardShrink and SoftShrink activation functions to support the ONNX Shrink operator. The implementation includes base tensor operations, neural network modules, configuration structs, display implementations, comprehensive tests, and documentation updates.

Changes:

• Added hard_shrink and soft_shrink activation functions in burn-tensor
• Created HardShrink and SoftShrink neural network modules with configuration support
• Integrated both activations into the activation wrapper system
• Updated documentation in the Burn book

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 8 comments.

Show a summary per file

File	Description
crates/burn-tensor/src/tensor/activation/base.rs	Implements base hard_shrink and soft_shrink activation functions with LaTeX documentation
crates/burn-nn/src/activation/hardshrink.rs	Defines HardShrink module, configuration, display, and tests
crates/burn-nn/src/activation/softshrink.rs	Defines SoftShrink module with separate lambd/bias parameters for ONNX compatibility
crates/burn-nn/src/activation/mod.rs	Exports the new activation modules
crates/burn-nn/src/lib.rs	Adds public exports for hardshrink and softshrink
crates/burn-nn/src/activation/activation_wrapper.rs	Integrates both activations into the activation enum with From implementations and tests
burn-book/src/building-blocks/tensor.md	Documents the new activation functions in the API reference
burn-book/src/building-blocks/module.md	Lists the new modules in the built-in modules table

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[antimora]

Nice work adding these activations - the implementations are correct and the integration into the Activation/ActivationConfig wrapper is clean. A few things to address though, see inline comments.

[antimora]

Minor fixes are required. Please see inlined comments.

[laggui]

Thanks for contributing!

If I understand correctly, I think ONNX Shrink is just a generalized implementation that handles both soft and hard shrink.

For bias = 0 it acts like HardShrink, and bias = lambda, it is like SoftShrink.

So if we add both hard_shrink and soft_shrink, we don't need soft shrink to have a separate bias no?

And for any other bias values, we could have a general shrink (reused in soft_shrink possibly).

[laggui]

@antimora looks like I started the review at the same time 😄 comments are pretty similar, except for the soft_shrink vs general shrink behavior. Lmk if my assessment is correct here.

[aditya0by0]

Thanks for contributing!

If I understand correctly, I think ONNX Shrink is just a generalized implementation that handles both soft and hard shrink.

For bias = 0 it acts like HardShrink, and bias = lambda, it is like SoftShrink.

So if we add both hard_shrink and soft_shrink, we don't need soft shrink to have a separate bias no?

And for any other bias values, we could have a general shrink (reused in soft_shrink possibly).

If I understand it correctly, we agree to add another ActivationConfig/Activation for general Shrink operation which internally calls burn::tensor::activation::soft_shrink with given bias value.

For Soft shink, we remove the bias parameter and the burn::tensor::activation::soft_shrink will be called with bias=lambda

[antimora]

@aditya0by0 Yeah that's exactly right. To spell it out:

• SoftShrink (nn module): just lambda, no separate bias. Calls soft_shrink(input, lambda, lambda) under the hood.
• HardShrink (nn module): stays as-is with just lambda.
• Shrink (nn module, new): has both lambda and bias params for the general ONNX case. Calls soft_shrink(input, lambda, bias) internally. This covers the arbitrary bias values that don't map to either hard or soft shrink.

At the tensor activation level, soft_shrink(tensor, lambda, bias) stays as the general implementation (maybe rename to just shrink?). hard_shrink can stay separate since its implementation is simpler (no sign/subtraction needed).

[laggui]

Yep, agreed with the above.

At the tensor activation level, soft_shrink(tensor, lambda, bias) stays as the general implementation (maybe rename to just shrink?).

Yeah we should probably rename to just shrink(tensor, lambda, bias), and to match the module struct we could have a shorthand soft_shrink(tensor, lambda) that just calls shrink(tensor, lambda, lambda) essentially?

[Copilot]

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated 1 comment.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

The table entries are not in alphabetical order. "Shrink" should be placed before "SoftShrink" to maintain alphabetical ordering in the table. The correct order should be: Selu, Shrink, SoftShrink, Softsign.

Suggested change

	| `SoftShrink` | `nn.Softshrink` |
	| `Softsign` | `nn.Softsign` |
	| `Shrink` | _No direct equivalent_ |
	| `Shrink` | _No direct equivalent_ |
	| `SoftShrink` | `nn.Softshrink` |
	| `Softsign` | `nn.Softsign` |

[aditya0by0]

I think the activations which don't have a pytorch equivalent goes in the bottom irrespective of alphabetical order.

[antimora]

All feedback addressed, looks good now. The three-tier split (HardShrink / SoftShrink / Shrink) is clean, soft_shrink properly delegates to shrink(tensor, lambda, lambda), LaTeX is fixed, unnecessary clone removed, file naming is consistent. Thanks for the quick turnaround!