v_diffusion.py

from shark.shark_inference import SharkInference
from shark.shark_downloader import download_torch_model
import numpy as np


mlir_model, func_name, inputs, golden_out = download_torch_model("v_diffusion")

shark_module = SharkInference(
    mlir_model, func_name, device="cpu", mlir_dialect="linalg"
)
shark_module.compile()
result = shark_module.forward(inputs)
print("The obtained result via shark is: ", result)
print("The golden result is:", golden_out)

tuple_gold = golden_out[0].reshape(2,3,256,256)

np.testing.assert_allclose(tuple_gold, result, rtol=1e-02, atol=1e-03)