aten._log_softmax.default failed with shape [19, 256008]

[swimdi]

aten._softmax.default failed with this input variation

aten__log_softmax_default_blocklist = [["Tensor<[19, 256008]> self = ?", "int dim = 1", "bool half_to_float = False"]]

Reproduce step is:

import torch
import ttnn

device = ttnn.open_device(device_id=0)
input = torch.rand([19, 256008])
ttnn_from_torch = ttnn.from_torch(input, device = device, layout = ttnn.TILE_LAYOUT, dtype = ttnn.bfloat16)
ttnn_softmax = ttnn.softmax(ttnn_from_torch, 1, numeric_stable = True)
ttnn_log = ttnn.log(ttnn_softmax)
ttnn_to_torch = ttnn.to_torch(ttnn_log, dtype = torch.bfloat16)

# Statically allocated circular buffers on core range [(x=0,y=0) - (x=7,y=7)] grow to 49277728 B which is beyond max L1 size of 1499136 B