Update memory config when using view op with height sharded tensors

This fixes a problem with ResNet50 (#17247) where the output memory
config was not being updated to match the new output width.

ttnn/cpp/ttnn/operations/experimental/reshape/view.cpp

@@ -13,18 +13,32 @@
 
 namespace ttnn::operations::experimental::reshape {
 
+static MemoryConfig infer_output_memory_config(
+    const MemoryConfig& input_memory_config, const ttnn::SimpleShape& output_logical_shape) {
+    if (input_memory_config.memory_layout == TensorMemoryLayout::HEIGHT_SHARDED) {
+        auto shard_spec = input_memory_config.shard_spec.value();
+        shard_spec.shape[1] = output_logical_shape[-1];  // update output shard to match new shard width
+        return MemoryConfig{input_memory_config.memory_layout, input_memory_config.buffer_type, shard_spec};
+    } else {
+        return input_memory_config;
+    }
+}
+
 Tensor tensor_reshape(
     const Tensor& input_tensor, const ttnn::SimpleShape& new_logical_shape, const ttnn::SimpleShape& new_padded_shape) {
     ZoneScoped;
     GraphTracker::instance().track_function_start("Tensor::reshape", input_tensor, new_logical_shape, new_padded_shape);
+
+    const auto output_memory_config = infer_output_memory_config(input_tensor.memory_config(), new_logical_shape);
     auto new_spec = ttnn::TensorSpec(
         new_logical_shape,
         TensorLayout::fromPaddedShape(
             input_tensor.get_dtype(),
             input_tensor.get_tensor_spec().page_config(),
-            input_tensor.memory_config(),
+            output_memory_config,
             new_logical_shape,
             new_padded_shape));
+
     auto output = std::visit(
         [&input_tensor, &new_spec, &new_logical_shape, &new_padded_shape](auto&& storage) -> Tensor {
             using T = std::decay_t<decltype(storage)>;