import numpy as np
from onnx import helper, TensorProto, save_model
from onnx.reference import ReferenceEvaluator

# Define inputs
x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [3, 4])
axis = helper.make_tensor_value_info('axis', TensorProto.INT32, [])

# Define outputs  
y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [3, 4])

# Create the CumProd node
cumprod_node = helper.make_node(
    'CumProd',
    inputs=['x', 'axis'],
    outputs=['y'],
    name='cumprod_node'
)

# Create the graph
graph = helper.make_graph(
    nodes=[cumprod_node],
    name='CumProdExample',
    inputs=[x, axis],
    outputs=[y]
)

# Create the model
model = helper.make_model(graph, producer_name='cumprod-example')
model.opset_import[0].version = 23
model.ir_version = 10

# Save the model
save_model(model, "temp_cumprod.onnx")

# Create reference evaluator for the model
ref_evaluator = ReferenceEvaluator("temp_cumprod.onnx")

# Prepare inputs
x_input = np.array([[1.0, 2.0, 3.0, 4.0],
                    [2.0, 2.0, 2.0, 2.0],
                    [1.0, 3.0, 1.0, 3.0]], dtype=np.float32)
axis_input = np.array(1, dtype=np.int32)  # Cumulative product along axis 1

# Run inference
outputs = ref_evaluator.run(None, {
    'x': x_input,
    'axis': axis_input
})

print(f"Input:\n{x_input}")
print(f"Axis: {axis_input}")
print(f"CumProd Output:\n{outputs[0]}")

# Compare with NumPy
numpy_result = np.cumprod(x_input, axis=1)
print(f"NumPy cumprod:\n{numpy_result}")
print(f"Results match: {np.allclose(outputs[0], numpy_result)}")