fix: transpose 1D channels_first input in TF depthwise_conv to support dilations by satishkc7 · Pull Request #22641 · keras-team/keras

satishkc7 · 2026-04-06T15:00:42Z

Summary

tf.nn.depthwise_conv2d does not support channels_first (NCHW) format with dilations on CPU. When DepthwiseConv1D is used with data_format='channels_first' and dilation_rate > 1, the TF backend's depthwise_conv function passed NCHW data directly to tf.nn.depthwise_conv2d, which fails at runtime with an Incompatible shapes error.

Root Cause

In keras/src/backend/tensorflow/nn.py, the 1D depthwise conv path for channels_first set spatial_start_dim = 2 and expanded the input as (N, C, 1, W) in NCHW format. Passing NCHW + dilations to tf.nn.depthwise_conv2d is not supported on CPU, causing the op to return incorrect results and ultimately an Incompatible shapes error when adding the bias.

Fix

For the 1D channels_first case, use the existing _transpose_spatial_inputs and _transpose_spatial_outputs helpers to transpose to channels_last before the convolution and back after, consistent with the pattern used by max_pool and average_pool in the same file.

Changes

keras/src/backend/tensorflow/nn.py: use _transpose_spatial_inputs/outputs helpers for 1D channels_first depthwise conv, compute in NHWC format
keras/src/layers/convolutional/depthwise_conv_test.py: add channels_first + dilation_rate=(2,) test cases to both the basic shape test and the correctness test

Fixes #22534

Contributor Agreement

I am a human, and not a bot.
I will be responsible for responding to review comments in a timely manner.
I will work with the maintainers to push this PR forward until submission.

tf.nn.depthwise_conv2d does not support channels_first (NCHW) format with dilations on CPU. When data_format='channels_first' and rank=1, transpose the input from (N, C, W) to (N, W, C) before expanding and running the convolution in NHWC format, then transpose the output back. Also adds test cases covering channels_first with dilation_rate > 1 in both the basic shape test and the correctness test. Fixes keras-team#22534

google-cla · 2026-04-06T15:00:55Z

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

gemini-code-assist

Code Review

This pull request addresses a limitation in the TensorFlow backend where 1D depthwise convolutions with dilations were unsupported on CPU when using the channels_first data format. The implementation transposes the data to channels_last for the operation and then reverts it. Feedback suggests using the existing internal helper functions _transpose_spatial_inputs and _transpose_spatial_outputs for these transpositions to ensure consistency with established patterns in the TensorFlow backend.

gemini-code-assist · 2026-04-06T15:06:15Z

keras/src/backend/tensorflow/nn.py

+        if data_format == "channels_first":
+            inputs = tf.transpose(inputs, perm=[0, 2, 1])


While this manual transpose works, the file already contains helper functions _transpose_spatial_inputs and _transpose_spatial_outputs (lines 188 and 206) that are used throughout the TensorFlow backend to handle exactly this kind of channels_first limitation. Using them would be more consistent with the rest of the file.

Suggested change

if data_format == "channels_first":

inputs = tf.transpose(inputs, perm=[0, 2, 1])

if data_format == "channels_first":

inputs = _transpose_spatial_inputs(inputs)

gemini-code-assist · 2026-04-06T15:06:15Z

keras/src/backend/tensorflow/nn.py

+        if data_format == "channels_first":
+            outputs = tf.transpose(outputs, perm=[0, 2, 1])


As with the input transpose, using the existing helper function _transpose_spatial_outputs would be more consistent with the repository's patterns.

Suggested change

if data_format == "channels_first":

outputs = tf.transpose(outputs, perm=[0, 2, 1])

if data_format == "channels_first":

outputs = _transpose_spatial_outputs(outputs)

codecov-commenter · 2026-04-06T15:06:18Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 83.26%. Comparing base (bfe3412) to head (ef327e6).
⚠️ Report is 1 commits behind head on master.

Additional details and impacted files

@@           Coverage Diff           @@
##           master   #22641   +/-   ##
=======================================
  Coverage   83.25%   83.26%           
=======================================
  Files         596      596           
  Lines       68244    68246    +2     
  Branches    10667    10668    +1     
=======================================
+ Hits        56816    56823    +7     
+ Misses       8638     8635    -3     
+ Partials     2790     2788    -2

Flag	Coverage Δ
keras	`83.07% <100.00%> (+<0.01%)`	⬆️
keras-jax	`59.62% <0.00%> (+<0.01%)`	⬆️
keras-numpy	`55.26% <0.00%> (+<0.01%)`	⬆️
keras-openvino	`59.03% <0.00%> (-0.01%)`	⬇️
keras-tensorflow	`60.96% <100.00%> (+0.01%)`	⬆️
keras-torch	`59.76% <0.00%> (+<0.01%)`	⬆️

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.

OpenVINO does not support channels_first for depthwise_conv and raises ValueError. Skip the new channels_first + dilation test cases when running on the OpenVINO backend.

hertschuh

Thanks for working on this fix!

hertschuh · 2026-04-13T18:27:06Z

keras/src/layers/convolutional/depthwise_conv_test.py

+        if data_format == "channels_first" and backend.backend() == "openvino":
+            self.skipTest(
+                "OpenVINO backend does not support channels_first for "
+                "depthwise_conv."
+            )


Can you add the test name to this list instead?
https://github.com/keras-team/keras/blob/master/keras/src/backend/openvino/excluded_concrete_tests.txt

This is how we track feature gaps in OpenVino

hertschuh · 2026-04-13T18:27:11Z

keras/src/layers/convolutional/depthwise_conv_test.py

+        if data_format == "channels_first" and backend.backend() == "openvino":
+            self.skipTest(
+                "OpenVINO backend does not support channels_first for "
+                "depthwise_conv."
+            )


Can you add the test name to this list instead?
https://github.com/keras-team/keras/blob/master/keras/src/backend/openvino/excluded_concrete_tests.txt

This is how we track feature gaps in OpenVino

hertschuh · 2026-04-13T18:28:24Z

keras/src/backend/tensorflow/nn.py

+        # `tf.nn.depthwise_conv2d` does not support `channels_first` with
+        # dilations on CPU. Transpose to `channels_last`, compute, and
+        # transpose back to avoid the limitation.
+        if data_format == "channels_first":
+            inputs = _transpose_spatial_inputs(inputs)


Can we add a check to only do this on CPU? I would assume transposing is slower than the current GPU / TPU implementations that currently work with "channels_first".

…cluded list

google-ml-butler bot added the size:S label Apr 6, 2026

google-ml-butler bot assigned gbaned Apr 6, 2026

gemini-code-assist bot reviewed Apr 6, 2026

View reviewed changes

satishkc7 added 2 commits April 6, 2026 11:11

fix: use _transpose_spatial_inputs/outputs helpers in depthwise_conv

138d321

fix: skip channels_first depthwise_conv tests on OpenVINO backend

ef327e6

OpenVINO does not support channels_first for depthwise_conv and raises ValueError. Skip the new channels_first + dilation test cases when running on the OpenVINO backend.

hertschuh added kokoro:force-run and removed kokoro:force-run labels Apr 6, 2026

keerthanakadiri added the stat:awaiting response from contributor label Apr 13, 2026

hertschuh reviewed Apr 13, 2026

View reviewed changes

fix: limit depthwise_conv transpose to CPU, move OpenVINO skips to ex…

c36ea35

…cluded list

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

fix: transpose 1D channels_first input in TF depthwise_conv to support dilations#22641

fix: transpose 1D channels_first input in TF depthwise_conv to support dilations#22641
satishkc7 wants to merge 4 commits intokeras-team:masterfrom
satishkc7:fix/flip-axis-list-tensorflow

satishkc7 commented Apr 6, 2026 •

edited

Loading

Uh oh!

google-cla bot commented Apr 6, 2026

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot Apr 6, 2026

Uh oh!

gemini-code-assist bot Apr 6, 2026

Uh oh!

codecov-commenter commented Apr 6, 2026 •

edited

Loading

Uh oh!

hertschuh left a comment

Uh oh!

hertschuh Apr 13, 2026

Uh oh!

hertschuh Apr 13, 2026

Uh oh!

hertschuh Apr 13, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

		if data_format == "channels_first":
		inputs = tf.transpose(inputs, perm=[0, 2, 1])

		if data_format == "channels_first":
		outputs = tf.transpose(outputs, perm=[0, 2, 1])

Conversation

satishkc7 commented Apr 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Root Cause

Fix

Changes

Contributor Agreement

Uh oh!

google-cla bot commented Apr 6, 2026

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Apr 6, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Apr 6, 2026

Choose a reason for hiding this comment

Uh oh!

codecov-commenter commented Apr 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

hertschuh left a comment

Choose a reason for hiding this comment

Uh oh!

hertschuh Apr 13, 2026

Choose a reason for hiding this comment

Uh oh!

hertschuh Apr 13, 2026

Choose a reason for hiding this comment

Uh oh!

hertschuh Apr 13, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

satishkc7 commented Apr 6, 2026 •

edited

Loading

codecov-commenter commented Apr 6, 2026 •

edited

Loading