import torch  # Not used but initializing the CUDA execution provider fails without it

import numpy as np
import onnx
from onnx.onnx_pb import TensorProto
import onnxruntime


num_cols = 1
num_rows = 3
operator = "Min"  # or "Max"

input = onnx.helper.make_tensor_value_info("input", TensorProto.FLOAT, ["N", num_cols])
output = onnx.helper.make_tensor_value_info("output", TensorProto.FLOAT, ["N", num_cols])

bound_const = onnx.helper.make_node(
        "Constant",
        inputs=[],
        outputs=["bound_const"],
        value=onnx.numpy_helper.from_array(np.array([10.0] * num_cols, dtype=np.float32)))

operator_node = onnx.helper.make_node(
        operator,
        inputs=["input", "bound_const"],
        outputs=["output"],
        )

graph_def = onnx.helper.make_graph(
        nodes=[bound_const, operator_node],
        name="test-model",
        inputs=[input],
        outputs=[output])

opset_import = onnx.helper.make_opsetid("", 17)

model_def = onnx.helper.make_model(
        graph_def,
        opset_imports=[opset_import],
        producer_name="test")

onnx.checker.check_model(model_def, full_check=True)

model_path = 'test_operator.onnx'
onnx.save(model_def, model_path)

input = np.random.randn(num_rows, num_cols).astype(np.float32)
input[0, :] = np.nan

cpu_session = onnxruntime.InferenceSession(model_path, providers=["CPUExecutionProvider"])
output = cpu_session.run(["output"], {"input": input})
print("CPU session output:")
print(output[0])

gpu_session = onnxruntime.InferenceSession(model_path, providers=["CUDAExecutionProvider"])
output = gpu_session.run(["output"], {"input": input})
print("GPU session output:")
print(output[0])