Error when compiling Bloom.py

[amalbasaTT]

Bloom compilation causes FATAL error:

File "/home/ubuntu/tt-metal/python_env/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
return forward_call(*args, **kwargs)
File "<eval_with_key>.16", line 57, in forward
ttnn_expand = ttnn_decorators_ttnn_expand(ttnn_to_device_4, [1, 1, 32, 32]); ttnn_to_device_4 = None
File "/home/ubuntu/tt-metal/ttnn/ttnn/decorators.py", line 333, in call
return self.function(*function_args, **function_kwargs)
RuntimeError: TT_FATAL @ /home/ubuntu/tt-metal/ttnn/cpp/ttnn/operations/data_movement/tilize/device/tilize_op.cpp:26: input_tensor_a.get_dtype() == DataType::BFLOAT16 or input_tensor_a.get_dtype() == DataType::FLOAT32
info:
data type must be bfloat16 or float32

To reproduce:

1. Run command:

pytest /home/ubuntu/pytorch2.0_ttnn/tests/models/bloom/test_bloom.py --gen_op_accuracy_tests

[kevinwuTT]

Here are the minimized versions that reproduce the same error.

import torch
import ttnn

arg295_1 = torch.randint(0, 10, (32, 1), dtype=torch.int64)

with ttnn.manage_device(device_id=0) as device:
    ttnn_from_torch_4 = ttnn.from_torch(arg295_1, layout = ttnn.ROW_MAJOR_LAYOUT, dtype = ttnn.uint32)
    ttnn_reshape_6 = ttnn.reshape(ttnn_from_torch_4, (1, 1, 32))
    ttnn_from_device_9 = ttnn.from_device(ttnn_reshape_6)
    ttnn_to_layout_9 = ttnn.to_layout(ttnn_from_device_9, ttnn.ROW_MAJOR_LAYOUT)
    ttnn_reshape_7 = ttnn.reshape(ttnn_to_layout_9, (1, 1, 1, 32))
    ttnn_from_device_10 = ttnn.from_device(ttnn_reshape_7)
    ttnn_to_layout_10 = ttnn.to_layout(ttnn_from_device_10, ttnn.TILE_LAYOUT)
    ttnn_to_device_4 = ttnn.to_device(ttnn_to_layout_10, device = device)
    ttnn_to_layout_10 = ttnn.to_layout(ttnn_from_device_10, ttnn.TILE_LAYOUT)
    ttnn_to_device_4 = ttnn.to_device(ttnn_to_layout_10, device = device)
    ttnn_expand = ttnn.expand(ttnn_to_device_4, [1, 1, 32, 32])

Minimized even further:

import torch
import ttnn

arg = torch.randint(0, 10, (1, 1, 1, 32), dtype=torch.int64)

with ttnn.manage_device(device_id=0) as device:
    ttnn_from_torch_4 = ttnn.from_torch(arg, layout = ttnn.TILE_LAYOUT, dtype = ttnn.uint32, device = device)
    ttnn_expand = ttnn.expand(ttnn_from_torch_4, [1, 1, 32, 32])

I think it comes down to whether ttnn.expand supports integer tensors anymore. If not, then we will fix this in pytorch2.0.