# python test_layer_norm.py
import torch

def forward(input):
    normalized_shape = (4, )
    weight = torch.ones(4, device="cuda")
    bias = torch.ones(4, device="cuda")
    eps = 0.1
    output = torch.layer_norm(
        input, normalized_shape, weight, bias, eps, torch.backends.cudnn.enabled
    )
    return output


x = torch.tensor([[1.0, 2.0, 3.0, 4.0],
              [2.0, 4.0, 6.0, 8.0]], device="cuda")

# no error
forward_compiled = torch.compile(forward, backend="aot_eager_decomp_partition")
forward_compiled(x.to(torch.bfloat16))

# error
forward_compiled = torch.compile(forward, backend="aot_eager")
forward_compiled(x.to(torch.bfloat16))

# error
# forward(x.to(torch.bfloat16))

RuntimeError: expected scalar type BFloat16 but found Float

While executing %native_layer_norm : [num_users=1] = call_function[target=torch.ops.aten.native_layer_norm.default](args = (%arg0_1, [4], %ones, %ones_1, 0.1), kwargs = {})
GraphModule: class <lambda>(torch.nn.Module):
    def forward(self, arg0_1: "bf16[2, 4][4, 1]"):
         # File: /data/users/weif/pytorch/test_layer_norm.py:5 in forward, code: weight = torch.ones(4, device="cuda")
        ones: "f32[4][1]" = torch.ops.aten.ones.default([4], device = device(type='cuda'), pin_memory = False)

         # File: /data/users/weif/pytorch/test_layer_norm.py:6 in forward, code: bias = torch.ones(4, device="cuda")
        ones_1: "f32[4][1]" = torch.ops.aten.ones.default([4], device = device(type='cuda'), pin_memory = False)

         # File: /data/users/weif/pytorch/test_layer_norm.py:8 in forward, code: output = torch.layer_norm(
        native_layer_norm = torch.ops.aten.native_layer_norm.default(arg0_1, [4], ones, ones_1, 0.1);  arg0_1 = ones = ones_1 = None
        getitem: "bf16[2, 4][4, 1]" = native_layer_norm[0];  native_layer_norm = None
        return (getitem,)