diff --git a/backends/cortex_m/ops/cortex_m_ops_common.h b/backends/cortex_m/ops/cortex_m_ops_common.h index bcaed0a1bc7..d6086f0e9c1 100644 --- a/backends/cortex_m/ops/cortex_m_ops_common.h +++ b/backends/cortex_m/ops/cortex_m_ops_common.h @@ -149,7 +149,10 @@ inline bool is_channels_last_tensor(const Tensor& tensor) { return tensor.dim_order() == channels_last_order; } -inline bool is_channel_broadcast(const Tensor& tensor1, const Tensor& tensor2) { +inline bool is_channel_broadcast( + const Tensor& tensor1, + const Tensor& tensor2, + int64_t channel_dim) { if (tensor1.dim() != tensor2.dim()) { return false; } @@ -158,16 +161,22 @@ inline bool is_channel_broadcast(const Tensor& tensor1, const Tensor& tensor2) { return false; } - if (tensor1.size(1) != tensor2.size(1)) { + if (tensor1.size(channel_dim) != tensor2.size(channel_dim)) { return false; } - const bool tensor1_channels_only = tensor1.numel() == tensor1.size(1); - const bool tensor2_channels_only = tensor2.numel() == tensor2.size(1); + const bool tensor1_channels_only = + tensor1.numel() == tensor1.size(channel_dim); + const bool tensor2_channels_only = + tensor2.numel() == tensor2.size(channel_dim); return tensor1_channels_only || tensor2_channels_only; } +inline bool is_channel_broadcast(const Tensor& tensor1, const Tensor& tensor2) { + return is_channel_broadcast(tensor1, tensor2, 1); +} + inline bool check_int32_within_range( KernelRuntimeContext& context, const char* op_name, diff --git a/backends/cortex_m/ops/op_quantized_add.cpp b/backends/cortex_m/ops/op_quantized_add.cpp index f93bb6c1be9..43ddd2f86de 100644 --- a/backends/cortex_m/ops/op_quantized_add.cpp +++ b/backends/cortex_m/ops/op_quantized_add.cpp @@ -13,8 +13,7 @@ namespace cortex_m { namespace native { using KernelRuntimeContext = torch::executor::KernelRuntimeContext; -// cppcheck-suppress unusedFunction -Tensor& quantized_add_out( +static Tensor& quantized_add_out_impl( KernelRuntimeContext& context, const Tensor& input1_int8, const int64_t input1_zero_point, @@ -29,16 +28,51 @@ Tensor& quantized_add_out( const int64_t output_shift, const int64_t activation_min, const int64_t activation_max, + ActivationLayout layout, + const char* op_name, Tensor& out) { - // Validate tensor types and dim order - bool channel_broadcast = is_channel_broadcast(input1_int8, input2_int8); + const int64_t channel_dim = layout == ActivationLayout::NHWCLogical ? 3 : 1; + bool channel_broadcast = + is_channel_broadcast(input1_int8, input2_int8, channel_dim); validate_cmsis_nn_tensor_requirements( input1_int8, input2_int8, out, ScalarType::Char, - /*require_channels_last=*/channel_broadcast, + /*require_channels_last=*/ + channel_broadcast && layout == ActivationLayout::NCHWLogical, /*require_same_sizes=*/!channel_broadcast); + if (layout == ActivationLayout::NHWCLogical) { + ET_CHECK_MSG( + input1_int8.dim() == 4 && input2_int8.dim() == 4 && out.dim() == 4, + "%s: tensors must be 4-D", + op_name); + ET_CHECK_MSG( + executorch::runtime::is_contiguous_dim_order( + input1_int8.dim_order().data(), input1_int8.dim_order().size()) && + executorch::runtime::is_contiguous_dim_order( + input2_int8.dim_order().data(), + input2_int8.dim_order().size()) && + executorch::runtime::is_contiguous_dim_order( + out.dim_order().data(), out.dim_order().size()), + "%s: tensors must use contiguous dimension order", + op_name); + } else if (channel_broadcast) { + ET_CHECK_MSG( + is_channels_last_tensor(input1_int8) && + is_channels_last_tensor(input2_int8) && + is_channels_last_tensor(out), + "%s: channel-broadcast tensors must use channels-last dimension order", + op_name); + } + if (channel_broadcast) { + const Tensor& full_input = + input1_int8.numel() > input2_int8.numel() ? input1_int8 : input2_int8; + ET_CHECK_MSG( + out.sizes() == full_input.sizes(), + "%s: output must have the broadcast result shape", + op_name); + } // Validate quantization parameters validate_quantization_params( @@ -54,7 +88,8 @@ Tensor& quantized_add_out( ET_LOG( Debug, - "quantized_add_out: input1_int8.sizes() = %zu", + "%s: input1_int8.sizes() = %zu", + op_name, input1_int8.sizes().size()); int32_t zp1 = static_cast(input1_zero_point); @@ -101,7 +136,7 @@ Tensor& quantized_add_out( std::swap(input1_shift_val, input2_shift_val); std::swap(input1_ptr, input2_ptr); } - adds_per_loop = input1_int8.size(1); + adds_per_loop = input1_int8.size(channel_dim); } else { adds_per_loop = out.numel(); } @@ -130,7 +165,8 @@ Tensor& quantized_add_out( if (status != ARM_CMSIS_NN_SUCCESS) { ET_LOG( Error, - "quantized_add_out: arm_elementwise_add_s8 failed with status [%d]", + "%s: arm_elementwise_add_s8 failed with status [%d]", + op_name, status); context.fail(Error::Internal); // Fail the execution context @@ -139,10 +175,85 @@ Tensor& quantized_add_out( } ET_LOG( Debug, - "quantized_add_out: Successfully completed with AoT-computed parameters!"); + "%s: Successfully completed with AoT-computed parameters!", + op_name); return out; } +// cppcheck-suppress unusedFunction +Tensor& quantized_add_out( + KernelRuntimeContext& context, + const Tensor& input1_int8, + const int64_t input1_zero_point, + const int64_t input1_multiplier, + const int64_t input1_shift, + const Tensor& input2_int8, + const int64_t input2_zero_point, + const int64_t input2_multiplier, + const int64_t input2_shift, + const int64_t output_zero_point, + const int64_t output_multiplier, + const int64_t output_shift, + const int64_t activation_min, + const int64_t activation_max, + Tensor& out) { + return quantized_add_out_impl( + context, + input1_int8, + input1_zero_point, + input1_multiplier, + input1_shift, + input2_int8, + input2_zero_point, + input2_multiplier, + input2_shift, + output_zero_point, + output_multiplier, + output_shift, + activation_min, + activation_max, + ActivationLayout::NCHWLogical, + "quantized_add_out", + out); +} + +// cppcheck-suppress unusedFunction +Tensor& quantized_add_nhwc_out( + KernelRuntimeContext& context, + const Tensor& input1_int8, + const int64_t input1_zero_point, + const int64_t input1_multiplier, + const int64_t input1_shift, + const Tensor& input2_int8, + const int64_t input2_zero_point, + const int64_t input2_multiplier, + const int64_t input2_shift, + const int64_t output_zero_point, + const int64_t output_multiplier, + const int64_t output_shift, + const int64_t activation_min, + const int64_t activation_max, + Tensor& out) { + return quantized_add_out_impl( + context, + input1_int8, + input1_zero_point, + input1_multiplier, + input1_shift, + input2_int8, + input2_zero_point, + input2_multiplier, + input2_shift, + output_zero_point, + output_multiplier, + output_shift, + activation_min, + activation_max, + ActivationLayout::NHWCLogical, + "quantized_add_nhwc_out", + out); +} + } // namespace native } // namespace cortex_m diff --git a/backends/cortex_m/ops/op_quantized_mul.cpp b/backends/cortex_m/ops/op_quantized_mul.cpp index 93ce2303d64..5b5f660b7d1 100644 --- a/backends/cortex_m/ops/op_quantized_mul.cpp +++ b/backends/cortex_m/ops/op_quantized_mul.cpp @@ -18,8 +18,7 @@ constexpr int32_t kInt8ActivationMax = std::numeric_limits::max(); using KernelRuntimeContext = torch::executor::KernelRuntimeContext; -// cppcheck-suppress unusedFunction -Tensor& quantized_mul_out( +static Tensor& quantized_mul_out_impl( KernelRuntimeContext& context, const Tensor& input1_int8, const int64_t input1_zero_point, @@ -28,17 +27,51 @@ Tensor& quantized_mul_out( const int64_t output_zero_point, const int64_t output_multiplier, const int64_t output_shift, + ActivationLayout layout, + const char* op_name, Tensor& out) { - // Validate tensor types and quantization parameters - - bool channel_broadcast = is_channel_broadcast(input1_int8, input2_int8); + const int64_t channel_dim = layout == ActivationLayout::NHWCLogical ? 3 : 1; + bool channel_broadcast = + is_channel_broadcast(input1_int8, input2_int8, channel_dim); validate_cmsis_nn_tensor_requirements( input1_int8, input2_int8, out, ScalarType::Char, - /*require_channels_last=*/channel_broadcast, + /*require_channels_last=*/ + channel_broadcast && layout == ActivationLayout::NCHWLogical, /*require_same_sizes=*/!channel_broadcast); + if (layout == ActivationLayout::NHWCLogical) { + ET_CHECK_MSG( + input1_int8.dim() == 4 && input2_int8.dim() == 4 && out.dim() == 4, + "%s: tensors must be 4-D", + op_name); + ET_CHECK_MSG( + executorch::runtime::is_contiguous_dim_order( + input1_int8.dim_order().data(), input1_int8.dim_order().size()) && + executorch::runtime::is_contiguous_dim_order( + input2_int8.dim_order().data(), + input2_int8.dim_order().size()) && + executorch::runtime::is_contiguous_dim_order( + out.dim_order().data(), out.dim_order().size()), + "%s: tensors must use contiguous dimension order", + op_name); + } else if (channel_broadcast) { + ET_CHECK_MSG( + is_channels_last_tensor(input1_int8) && + is_channels_last_tensor(input2_int8) && + is_channels_last_tensor(out), + "%s: channel-broadcast tensors must use channels-last dimension order", + op_name); + } + if (channel_broadcast) { + const Tensor& full_input = + input1_int8.numel() > input2_int8.numel() ? input1_int8 : input2_int8; + ET_CHECK_MSG( + out.sizes() == full_input.sizes(), + "%s: output must have the broadcast result shape", + op_name); + } const int32_t kIdentityMultiplier(/*value=*/1); const int32_t kZeroShift(/*value=*/0); @@ -70,7 +103,7 @@ Tensor& quantized_mul_out( std::swap(input1_ptr, input2_ptr); } - muls_per_loop = input1_int8.size(1); + muls_per_loop = input1_int8.size(channel_dim); } else { muls_per_loop = out.numel(); } @@ -108,7 +141,8 @@ Tensor& quantized_mul_out( if (status != ARM_CMSIS_NN_SUCCESS) { ET_LOG( Error, - "quantized_mul_out: arm_elementwise_mul_s8 failed with status [%d]", + "%s: arm_elementwise_mul_s8 failed with status [%d]", + op_name, status); context.fail(Error::Internal); return out; @@ -117,5 +151,55 @@ Tensor& quantized_mul_out( return out; } +// cppcheck-suppress unusedFunction +Tensor& quantized_mul_out( + KernelRuntimeContext& context, + const Tensor& input1_int8, + const int64_t input1_zero_point, + const Tensor& input2_int8, + const int64_t input2_zero_point, + const int64_t output_zero_point, + const int64_t output_multiplier, + const int64_t output_shift, + Tensor& out) { + return quantized_mul_out_impl( + context, + input1_int8, + input1_zero_point, + input2_int8, + input2_zero_point, + output_zero_point, + output_multiplier, + output_shift, + ActivationLayout::NCHWLogical, + "quantized_mul_out", + out); +} + +// cppcheck-suppress unusedFunction +Tensor& quantized_mul_nhwc_out( + KernelRuntimeContext& context, + const Tensor& input1_int8, + const int64_t input1_zero_point, + const Tensor& input2_int8, + const int64_t input2_zero_point, + const int64_t output_zero_point, + const int64_t output_multiplier, + const int64_t output_shift, + Tensor& out) { + return quantized_mul_out_impl( + context, + input1_int8, + input1_zero_point, + input2_int8, + input2_zero_point, + output_zero_point, + output_multiplier, + output_shift, + ActivationLayout::NHWCLogical, + "quantized_mul_nhwc_out", + out); +} + } // namespace native } // namespace cortex_m diff --git a/backends/cortex_m/ops/operators.py b/backends/cortex_m/ops/operators.py index 79c7dddcc30..bcc78fafddd 100644 --- a/backends/cortex_m/ops/operators.py +++ b/backends/cortex_m/ops/operators.py @@ -137,6 +137,23 @@ def dequantize_per_tensor_impl( "*, Tensor(a!) out) -> Tensor(a!)" ) +lib.define( + "quantized_add_nhwc(" + "Tensor self, int self_zero_point, int self_multiplier, int self_shift, " + "Tensor other, int other_zero_point, int other_multiplier, int other_shift, " + "int output_zero_point, int output_multiplier, int output_shift, " + "int activation_min, int activation_max) -> Tensor" +) + +lib.define( + "quantized_add_nhwc.out(" + "Tensor self, int self_zero_point, int self_multiplier, int self_shift, " + "Tensor other, int other_zero_point, int other_multiplier, int other_shift, " + "int output_zero_point, int output_multiplier, int output_shift, " + "int activation_min, int activation_max, " + "*, Tensor(a!) out) -> Tensor(a!)" +) + @register_fake("cortex_m::quantized_add") # type: ignore[misc] def quantized_add_meta( @@ -199,6 +216,78 @@ def quantized_add_impl( return result +@register_fake("cortex_m::quantized_add_nhwc") # type: ignore[misc] +def quantized_add_nhwc_meta( + self: torch.Tensor, + self_zero_point: int, + self_multiplier: int, + self_shift: int, + other: torch.Tensor, + other_zero_point: int, + other_multiplier: int, + other_shift: int, + output_zero_point: int, + output_multiplier: int, + output_shift: int, + activation_min: int, + activation_max: int, +) -> torch.Tensor: + if self.dim() != 4 or other.dim() != 4: + raise RuntimeError("cortex_m.quantized_add_nhwc expects 4D inputs") + result = quantized_add_meta( + self.permute(0, 3, 1, 2), + self_zero_point, + self_multiplier, + self_shift, + other.permute(0, 3, 1, 2), + other_zero_point, + other_multiplier, + other_shift, + output_zero_point, + output_multiplier, + output_shift, + activation_min, + activation_max, + ) + return result.permute(0, 2, 3, 1).contiguous() + + +@impl(lib, "quantized_add_nhwc", "CompositeExplicitAutograd") # type: ignore[misc] +def quantized_add_nhwc_impl( + self: torch.Tensor, + self_zero_point: int, + self_multiplier: int, + self_shift: int, + other: torch.Tensor, + other_zero_point: int, + other_multiplier: int, + other_shift: int, + output_zero_point: int, + output_multiplier: int, + output_shift: int, + activation_min: int, + activation_max: int, +) -> torch.Tensor: + if self.dim() != 4 or other.dim() != 4: + raise RuntimeError("cortex_m.quantized_add_nhwc expects 4D inputs") + result = quantized_add_impl( + self.permute(0, 3, 1, 2), + self_zero_point, + self_multiplier, + self_shift, + other.permute(0, 3, 1, 2), + other_zero_point, + other_multiplier, + other_shift, + output_zero_point, + output_multiplier, + output_shift, + activation_min, + activation_max, + ) + return result.permute(0, 2, 3, 1).contiguous() + + # =================================================================== # QUANTIZED MUL OPERATION DEFINITION # =================================================================== @@ -216,6 +305,20 @@ def quantized_add_impl( "*, Tensor(a!) out) -> Tensor(a!)" ) +lib.define( + "quantized_mul_nhwc(" + "Tensor self, int self_zero_point, " + "Tensor other, int other_zero_point, " + "int output_zero_point, int output_multiplier, int output_shift) -> Tensor" +) +lib.define( + "quantized_mul_nhwc.out(" + "Tensor self, int self_zero_point, " + "Tensor other, int other_zero_point, " + "int output_zero_point, int output_multiplier, int output_shift, " + "*, Tensor(a!) out) -> Tensor(a!)" +) + @register_fake("cortex_m::quantized_mul") # type: ignore[misc] def quantized_mul_meta( @@ -264,6 +367,54 @@ def quantized_mul_impl( return result +@register_fake("cortex_m::quantized_mul_nhwc") # type: ignore[misc] +def quantized_mul_nhwc_meta( + self: torch.Tensor, + self_zero_point: int, + other: torch.Tensor, + other_zero_point: int, + output_zero_point: int, + output_multiplier: int, + output_shift: int, +) -> torch.Tensor: + if self.dim() != 4 or other.dim() != 4: + raise RuntimeError("cortex_m.quantized_mul_nhwc expects 4D inputs") + result = quantized_mul_meta( + self.permute(0, 3, 1, 2), + self_zero_point, + other.permute(0, 3, 1, 2), + other_zero_point, + output_zero_point, + output_multiplier, + output_shift, + ) + return result.permute(0, 2, 3, 1).contiguous() + + +@impl(lib, "quantized_mul_nhwc", "CompositeExplicitAutograd") # type: ignore[misc] +def quantized_mul_nhwc_impl( + self: torch.Tensor, + self_zero_point: int, + other: torch.Tensor, + other_zero_point: int, + output_zero_point: int, + output_multiplier: int, + output_shift: int, +) -> torch.Tensor: + if self.dim() != 4 or other.dim() != 4: + raise RuntimeError("cortex_m.quantized_mul_nhwc expects 4D inputs") + result = quantized_mul_impl( + self.permute(0, 3, 1, 2), + self_zero_point, + other.permute(0, 3, 1, 2), + other_zero_point, + output_zero_point, + output_multiplier, + output_shift, + ) + return result.permute(0, 2, 3, 1).contiguous() + + # =================================================================== # QUANTIZED DIV OPERATION DEFINITION # =================================================================== diff --git a/backends/cortex_m/ops/operators.yaml b/backends/cortex_m/ops/operators.yaml index e91aaca3569..43b75c784c4 100644 --- a/backends/cortex_m/ops/operators.yaml +++ b/backends/cortex_m/ops/operators.yaml @@ -23,12 +23,24 @@ - arg_meta: null kernel_name: cortex_m::quantized_add_out +- func: cortex_m::quantized_add_nhwc.out(Tensor self, int self_zero_point, int self_multiplier, int self_shift, Tensor other, int other_zero_point, int other_multiplier, int other_shift, int output_zero_point, int output_multiplier, int output_shift, int activation_min, int activation_max, *, Tensor(a!) out) -> Tensor(a!) + variants: function + kernels: + - arg_meta: null + kernel_name: cortex_m::quantized_add_nhwc_out + - func: cortex_m::quantized_mul.out(Tensor self, int self_zero_point, Tensor other, int other_zero_point, int output_zero_point, int output_multiplier, int output_shift, *, Tensor(a!) out) -> Tensor(a!) variants: function kernels: - arg_meta: null kernel_name: cortex_m::quantized_mul_out +- func: cortex_m::quantized_mul_nhwc.out(Tensor self, int self_zero_point, Tensor other, int other_zero_point, int output_zero_point, int output_multiplier, int output_shift, *, Tensor(a!) out) -> Tensor(a!) + variants: function + kernels: + - arg_meta: null + kernel_name: cortex_m::quantized_mul_nhwc_out + - func: cortex_m::quantized_div.out(Tensor self, int self_zero_point, Tensor other, int other_zero_point, int output_zero_point, float output_scale, *, Tensor(a!) out) -> Tensor(a!) variants: function kernels: diff --git a/backends/cortex_m/test/build_test_runner.sh b/backends/cortex_m/test/build_test_runner.sh index ad91eef264a..f1e301f1bcd 100755 --- a/backends/cortex_m/test/build_test_runner.sh +++ b/backends/cortex_m/test/build_test_runner.sh @@ -56,8 +56,10 @@ ops_list=( cortex_m::quantize_per_tensor.out cortex_m::dequantize_per_tensor.out cortex_m::quantized_add.out + cortex_m::quantized_add_nhwc.out cortex_m::quantized_div.out cortex_m::quantized_mul.out + cortex_m::quantized_mul_nhwc.out cortex_m::quantized_activation.out cortex_m::minimum.out cortex_m::maximum.out diff --git a/backends/cortex_m/test/ops/test_explicit_nhwc_runtime.py b/backends/cortex_m/test/ops/test_explicit_nhwc_runtime.py index 36b92bbd9b5..aa2ad366829 100644 --- a/backends/cortex_m/test/ops/test_explicit_nhwc_runtime.py +++ b/backends/cortex_m/test/ops/test_explicit_nhwc_runtime.py @@ -197,6 +197,46 @@ def forward(self, x): ) +class AddNhwc(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("bias", _int8_values((1, 1, 1, 3))) + + def forward(self, x): + return torch.ops.cortex_m.quantized_add_nhwc.default( + x, + 0, + 1 << 30, + -1, + self.bias, + 0, + 1 << 30, + -1, + 0, + 1 << 30, + -1, + -128, + 127, + ) + + +class MulNhwc(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("bias", _int8_values((1, 1, 1, 3))) + + def forward(self, x): + return torch.ops.cortex_m.quantized_mul_nhwc.default( + x, + 0, + self.bias, + 0, + 0, + 1 << 30, + -1, + ) + + def test_conv2d_nhwc_runs_on_fvp(cortex_m_target): _run_on_fvp( Conv2dNhwc(), @@ -254,3 +294,23 @@ def test_pad_nhwc_runs_on_fvp_with_singleton_height(cortex_m_target): exir_ops.edge.cortex_m.pad_nhwc.default, cortex_m_target, ) + + +def test_channel_broadcast_add_nhwc_runs_on_fvp(cortex_m_target): + _run_on_fvp( + AddNhwc(), + _int8_values((1, 5, 7, 3)), + exir_ops.edge.cortex_m.quantized_add_nhwc.default, + cortex_m_target, + atol=1, + ) + + +def test_channel_broadcast_mul_nhwc_runs_on_fvp(cortex_m_target): + _run_on_fvp( + MulNhwc(), + _int8_values((1, 5, 7, 3)), + exir_ops.edge.cortex_m.quantized_mul_nhwc.default, + cortex_m_target, + atol=1, + ) diff --git a/backends/cortex_m/test/test_quantized_conv2d_layout.py b/backends/cortex_m/test/test_quantized_conv2d_layout.py index 891ea750686..7a99a372e92 100644 --- a/backends/cortex_m/test/test_quantized_conv2d_layout.py +++ b/backends/cortex_m/test/test_quantized_conv2d_layout.py @@ -104,6 +104,28 @@ def _run_max_pool2d(op, x): ) +def _run_add(op, x, bias): + return op( + x, + 0, + 1 << 30, + -1, + bias, + 0, + 1 << 30, + -1, + 0, + 1 << 30, + -1, + -128, + 127, + ) + + +def _run_mul(op, x, bias): + return op(x, 0, bias, 0, 0, 1 << 30, -1) + + def test_nhwc_conv2d_matches_legacy_layout(): torch.manual_seed(0) x = torch.randint(-8, 8, (1, 3, 8, 8), dtype=torch.int8) @@ -200,6 +222,42 @@ def test_nhwc_max_pool2d_matches_legacy_layout(): torch.testing.assert_close(explicit, legacy.permute(0, 2, 3, 1)) +def test_nhwc_channel_broadcast_add_matches_legacy_layout(): + x = torch.randint(-8, 8, (1, 4, 5, 7), dtype=torch.int8) + bias = torch.randint(-4, 4, (1, 4, 1, 1), dtype=torch.int8) + + legacy = _run_add( + torch.ops.cortex_m.quantized_add, + x.to(memory_format=torch.channels_last), + bias.to(memory_format=torch.channels_last), + ) + explicit = _run_add( + torch.ops.cortex_m.quantized_add_nhwc, + x.permute(0, 2, 3, 1).contiguous(), + bias.permute(0, 2, 3, 1).contiguous(), + ) + + torch.testing.assert_close(explicit, legacy.permute(0, 2, 3, 1)) + + +def test_nhwc_channel_broadcast_mul_matches_legacy_layout(): + x = torch.randint(-8, 8, (1, 4, 5, 7), dtype=torch.int8) + bias = torch.randint(-4, 4, (1, 4, 1, 1), dtype=torch.int8) + + legacy = _run_mul( + torch.ops.cortex_m.quantized_mul, + x.to(memory_format=torch.channels_last), + bias.to(memory_format=torch.channels_last), + ) + explicit = _run_mul( + torch.ops.cortex_m.quantized_mul_nhwc, + x.permute(0, 2, 3, 1).contiguous(), + bias.permute(0, 2, 3, 1).contiguous(), + ) + + torch.testing.assert_close(explicit, legacy.permute(0, 2, 3, 1)) + + def test_nhwc_conv2d_fake_shape_is_logical_nhwc(): with FakeTensorMode(): output = _run_conv2d(