[Bug] input_type bfloat16 lead to wrong predict results

[jikechao]

Actual behavior

What actually happened

Environment

Any environment details, such as: Operating System, TVM version, etc

Steps to reproduce

import tvm
import tvm.relay as relay
import numpy as np
from tensorflow import keras
from tensorflow.keras import layers, models


input_shape = (1, 12, 4, 6)
input_data = np.random.random(input_shape)
x = layers.Input(shape=input_shape[1:], dtype='bfloat16')

layer = keras.layers.Dropout(rate=0, seed=[1, 1])
#layer.set_weights(layer.get_weights())

y = layer(x)
model = models.Model(x, y)
#res_keras = model.predict(input_data)
res_keras = model(input_data)
shape_dict = {'input_1': input_shape}
mod, params = relay.frontend.from_keras(model, shape_dict)

with tvm.transform.PassContext(opt_level=3):
    model = relay.build_module.create_executor("vm", mod, tvm.cpu(0), 'llvm', params).evaluate()

res_tvm = model(tvm.nd.array(input_data.astype('bfloat16'))).numpy()

print('keras infer result:', res_keras)
print('tvm infer result:', res_tvm)
np.testing.assert_allclose(res_keras, res_tvm, atol=1e-3, rtol=1e-3)

Notice that "float32" can receive a correct inference result.