Skip to content

Add fold to keras.ops.nn#22283

Open
MarcosAsh wants to merge 2 commits intokeras-team:masterfrom
MarcosAsh:add-fold-op
Open

Add fold to keras.ops.nn#22283
MarcosAsh wants to merge 2 commits intokeras-team:masterfrom
MarcosAsh:add-fold-op

Conversation

@MarcosAsh
Copy link
Contributor

Summary

  • Adds keras.ops.fold, the inverse of keras.ops.unfold (im2col). Takes a 3D tensor of flattened patches (N, CkHkW, L) and reconstructs a 4D
    tensor (N, C, oH, oW) by summing overlapping patches.
  • Implements backends for torch (passthrough to tnn.fold), jax, tensorflow, numpy, and openvino (NotImplementedError, matching unfold).
  • Addresses part of Feature request: fold/unfold ops #20046

Test plan

  • Non overlapping roundtrip: fold(unfold(x), ...) == x when stride == kernel_size
  • Overlapping roundtrip: fold(unfold(x)) / fold(ones) == x
  • Multi channel non overlapping roundtrip
  • Dilation + padding roundtrip
  • Explicit known values
  • Input validation (must be 3D)

@gemini-code-assist
Copy link
Contributor

Summary of Changes

Hello @MarcosAsh, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request adds the fold operation to keras.ops.nn, providing the inverse functionality to unfold. This is achieved through implementations across multiple backends and a comprehensive suite of tests to ensure correctness and compatibility.

Highlights

  • Fold Implementation: This PR introduces the keras.ops.fold operation, which is the inverse of keras.ops.unfold (im2col), enabling the reconstruction of a 4D tensor from flattened patches.
  • Backend Support: The fold operation is implemented with backend support for Torch (passthrough to tnn.fold), JAX, TensorFlow, and NumPy. OpenVINO receives a NotImplementedError.
  • Comprehensive Testing: The PR includes a thorough test plan covering non-overlapping and overlapping roundtrips, multi-channel scenarios, dilation and padding, explicit known values, and input validation.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog
  • keras/api/_tf_keras/keras/ops/init.py
    • Imported the fold function from keras.src.ops.nn.
  • keras/api/_tf_keras/keras/ops/nn/init.py
    • Imported the fold function from keras.src.ops.nn.
  • keras/api/ops/init.py
    • Imported the fold function from keras.src.ops.nn.
  • keras/api/ops/nn/init.py
    • Imported the fold function from keras.src.ops.nn.
  • keras/src/backend/jax/nn.py
    • Implemented the fold operation for the JAX backend.
  • keras/src/backend/numpy/nn.py
    • Implemented the fold operation for the NumPy backend.
  • keras/src/backend/openvino/nn.py
    • Implemented the fold operation for the OpenVINO backend, raising a NotImplementedError.
  • keras/src/backend/tensorflow/nn.py
    • Implemented the fold operation for the TensorFlow backend.
  • keras/src/backend/torch/nn.py
    • Implemented the fold operation for the Torch backend, using tnn.fold.
  • keras/src/ops/nn.py
    • Added the Fold operation class and the fold function, including input validation and backend dispatch.
  • keras/src/ops/nn_test.py
    • Added comprehensive tests for the fold operation, covering various scenarios and input validations.
Activity
  • The pull request introduces a new fold operation to the Keras ops API.
  • It includes implementations for multiple backends (JAX, NumPy, TensorFlow, Torch) and a placeholder for OpenVINO.
  • The PR also adds a comprehensive test suite to ensure the correctness of the new operation.
Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help /gemini help Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution.

Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces the fold operation to keras.ops.nn, which serves as the inverse of unfold (im2col). The implementation includes support for JAX, TensorFlow, and NumPy backends, a passthrough for the Torch backend, and raises a NotImplementedError for the OpenVINO backend, which is consistent with the unfold operation. The changes are well-structured and include a comprehensive set of tests covering various scenarios like non-overlapping and overlapping roundtrips, multi-channel inputs, and different parameter configurations. My review includes one suggestion for the TensorFlow backend implementation to improve code consistency and maintainability by using a helper function, similar to how it's done in the JAX and NumPy backends.

Comment on lines +1641 to +1653
k = (
(kernel_size, kernel_size)
if isinstance(kernel_size, int)
else kernel_size
)
o = (
(output_size, output_size)
if isinstance(output_size, int)
else output_size
)
d = (dilation, dilation) if isinstance(dilation, int) else dilation
p = (padding, padding) if isinstance(padding, int) else padding
s = (stride, stride) if isinstance(stride, int) else stride
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

For consistency with the JAX and NumPy backend implementations of fold, consider using a _pair helper function to handle integer inputs for kernel_size, output_size, dilation, padding, and stride. This would reduce code duplication and improve readability, adhering to the principle of internal consistency mentioned in the style guide.

    def _pair(x):
        return (x, x) if isinstance(x, int) else x

    k = _pair(kernel_size)
    o = _pair(output_size)
    d = _pair(dilation)
    p = _pair(padding)
    s = _pair(stride)
References
  1. The style guide emphasizes internal consistency. While the rule on line 71 specifically mentions naming, its principle of internal consistency can be extended to implementation patterns. Using a _pair helper function here would align the TensorFlow backend's implementation with the JAX and NumPy backends for the same operation, improving code clarity and maintainability. (link)

@codecov-commenter
Copy link

codecov-commenter commented Feb 25, 2026

Codecov Report

❌ Patch coverage is 84.25197% with 20 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.74%. Comparing base (0ddf962) to head (a923dfa).
⚠️ Report is 22 commits behind head on master.

Files with missing lines Patch % Lines
keras/src/ops/nn.py 41.93% 17 Missing and 1 partial ⚠️
keras/api/_tf_keras/keras/ops/__init__.py 0.00% 1 Missing ⚠️
keras/api/_tf_keras/keras/ops/nn/__init__.py 0.00% 1 Missing ⚠️
Additional details and impacted files
@@             Coverage Diff             @@
##           master   #22283       +/-   ##
===========================================
+ Coverage   71.44%   82.74%   +11.29%     
===========================================
  Files         594      594               
  Lines       65029    65860      +831     
  Branches    10174    10278      +104     
===========================================
+ Hits        46461    54493     +8032     
+ Misses      16105     8730     -7375     
- Partials     2463     2637      +174     
Flag Coverage Δ
keras 82.56% <84.25%> (+11.24%) ⬆️
keras-jax 60.86% <34.64%> (-0.58%) ⬇️
keras-numpy 55.07% <35.43%> (-0.56%) ⬇️
keras-openvino 39.28% <12.59%> (?)
keras-tensorflow 62.12% <41.73%> (?)
keras-torch 60.93% <14.96%> (-0.62%) ⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ 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.

@MarcosAsh MarcosAsh force-pushed the add-fold-op branch 2 times, most recently from a476465 to 6245fec Compare February 25, 2026 09:44
Fix ruff formatting in tensorflow/nn.py and nn_test.py. Add tests
for tuple parameters, zero padding, non-square outputs, larger
batch/channel dimensions, and divisibility validation.
@MarcosAsh
Copy link
Contributor Author

Hi I did some changes I added tests for tuple parameters (non-square kernels, asymmetric padding/stride) and added tests for zero padding, non-square output sizes, larger batch/channel dimensions, and divisibility validation

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants