Add support for NCHW layout in stablehlo.composites for jax.image.resize and torch.nn.interpolate in nearest mode.

[copybara-service]

Add support for NCHW layout in stablehlo.composites for jax.image.resize and torch.nn.interpolate in nearest mode.

[NIV27e]

Great job on this PR! Adding support for the NCHW layout in stablehlo.composites for jax.image.resize and torch.nn.interpolate in nearest mode is a significant improvement.

Composite-Lowering.mlir
The is_nchw_op attribute addition is well done. This clearly differentiates between NCHW and NHWC operations. Just ensure all edge cases are covered.

Composite_Lowering_Patterns.td
The new constraints and patterns for NCHW are clearly defined and seem to work well. Make sure they’re thoroughly tested with various inputs.

Composite_Utils.cc and Composite_Utils.h
The removal of IsSupportedNchwUpsampleBlinear and related utilities makes sense if they’re no longer needed. Just double-check that all utility functions handle NCHW layout without performance issues.

Your comments are clear and helpful. If possible, simplify any complex logic to enhance readability. Consider profiling the changes to avoid any performance hits, especially in critical paths.

Questions

1. Are there specific test cases for verifying the new NCHW layout handling?
2. How do these changes affect workflows that rely on NHWC layout? Any backward compatibility issues?

Overall, this is solid work! Looking forward to seeing it merged.