@@ -1503,24 +1503,24 @@ def verify_embedlayernormalization(
15031503 )
15041504
15051505
1506- def create_reduce_test_parameters ():
1506+ def create_reduce_test_parameters_axes_attr ():
15071507 output = []
15081508 for value in [True , False ]:
1509- output .append (("ReduceMax" , value ))
1510- output .append (("ReduceMean" , value ))
1511- output .append (("ReduceMin" , value ))
1512- output .append (("ReduceProd" , value ))
1513- output .append (("ReduceSum" , value ))
1514- output .append (("ReduceSumSquare" , value ))
1515- output .append (("ReduceLogSum" , value ))
1516- output .append (("ReduceLogSumExp" , value ))
1517- output .append (("ReduceL1" , value ))
1518- output .append (("ReduceL2" , value ))
1509+ output .append (("ReduceMax" , value , 11 ))
1510+ output .append (("ReduceMean" , value , 13 ))
1511+ output .append (("ReduceMin" , value , 11 ))
1512+ output .append (("ReduceProd" , value , 13 ))
1513+ output .append (("ReduceSum" , value , 11 ))
1514+ output .append (("ReduceSumSquare" , value , 13 ))
1515+ output .append (("ReduceLogSum" , value , 13 ))
1516+ output .append (("ReduceLogSumExp" , value , 13 ))
1517+ output .append (("ReduceL1" , value , 13 ))
1518+ output .append (("ReduceL2" , value , 13 ))
15191519 return output
15201520
15211521
1522- @pytest .mark .parametrize ("func, dynamic" , create_reduce_test_parameters ())
1523- def test_all_reduce_funcs (func , dynamic ):
1522+ @pytest .mark .parametrize ("func, dynamic, opset " , create_reduce_test_parameters_axes_attr ())
1523+ def test_all_reduce_funcs_axes_attr (func , dynamic , opset ):
15241524 def verify_reduce_func (func , data , axis , keepdims ):
15251525 inshape = data .shape
15261526 outshape = np .sum (data , axis = axis , keepdims = keepdims == 1 ).shape
@@ -1549,7 +1549,7 @@ def verify_reduce_func(func, data, axis, keepdims):
15491549
15501550 inputs_dict = {"x" : data }
15511551 # Reduction ops accumulate arithmetic errors, so we use a higher tolerance.
1552- check_correctness (model , inputs_dict , opset = 11 , rtol = 1e-4 , atol = 1e-4 )
1552+ check_correctness (model , inputs_dict , opset = opset , rtol = 1e-4 , atol = 1e-4 )
15531553
15541554 for keepdims in [True , False ]:
15551555 verify_reduce_func (
@@ -1577,6 +1577,131 @@ def verify_reduce_func(func, data, axis, keepdims):
15771577 )
15781578
15791579
1580+ def create_reduce_test_parameters_axes_input ():
1581+ output = []
1582+ for dynamic in [True , False ]:
1583+ # TODO(@vacu9708): Enable the tests after implementing other reduce ops
1584+ # output.append(("ReduceMax", dynamic, 20))
1585+ # output.append(("ReduceMean", dynamic, 18))
1586+ # output.append(("ReduceMin", dynamic, 20))
1587+ # output.append(("ReduceProd", dynamic, 18))
1588+ # output.append(("ReduceSum", dynamic, 13))
1589+ # output.append(("ReduceSumSquare", dynamic, 18))
1590+ # output.append(("ReduceLogSum", dynamic, 18))
1591+ # output.append(("ReduceLogSumExp", dynamic, 18))
1592+ output .append (("ReduceL1" , dynamic , 18 ))
1593+ # output.append(("ReduceL2", dynamic, 18))
1594+ return output
1595+
1596+
1597+ @pytest .mark .parametrize ("func, dynamic, opset" , create_reduce_test_parameters_axes_input ())
1598+ def test_all_reduce_funcs_axes_input (func , dynamic , opset ):
1599+ def verify_reduce_func (func , data , axes , keepdims , noop_with_empty_axes ):
1600+ inshape = data .shape
1601+
1602+ inputs = ["x" ]
1603+ initializers = []
1604+
1605+ # Optional `axes` input
1606+ if axes is not None :
1607+ axes_name = "reduce_axes"
1608+ axes_np = np .asarray (axes , dtype = np .int64 )
1609+ axes_init = helper .make_tensor (
1610+ name = axes_name ,
1611+ data_type = TensorProto .INT64 ,
1612+ dims = axes_np .shape ,
1613+ vals = axes_np ,
1614+ )
1615+ initializers .append (axes_init )
1616+ inputs .append (axes_name )
1617+
1618+ # Determine input and output shapes
1619+ if not axes and not noop_with_empty_axes :
1620+ outshape = np .sum (data , axis = None , keepdims = keepdims ).shape
1621+ elif not axes and noop_with_empty_axes :
1622+ outshape = inshape
1623+ else :
1624+ outshape = np .sum (data , axis = axes , keepdims = keepdims ).shape
1625+
1626+ if dynamic :
1627+ in_list = ["?" ] * len (inshape )
1628+ out_list = ["?" ] * len (outshape )
1629+ else :
1630+ in_list = list (inshape )
1631+ out_list = list (outshape )
1632+
1633+ # Make a model node
1634+ node = helper .make_node (
1635+ func ,
1636+ inputs = inputs ,
1637+ outputs = ["y" ],
1638+ keepdims = keepdims ,
1639+ noop_with_empty_axes = noop_with_empty_axes ,
1640+ )
1641+
1642+ # Make a model graph and a model
1643+ graph = helper .make_graph (
1644+ [node ],
1645+ "reduce18_test" ,
1646+ inputs = [helper .make_tensor_value_info ("x" , TensorProto .FLOAT , in_list )],
1647+ initializer = initializers ,
1648+ outputs = [helper .make_tensor_value_info ("y" , TensorProto .FLOAT , out_list )],
1649+ )
1650+ model = helper .make_model (graph , producer_name = "reduce18_test" )
1651+
1652+ # Run TVM importer vs onnxruntime
1653+ inputs_dict = {"x" : data }
1654+ check_correctness (model , inputs_dict , opset = opset , rtol = 1e-4 , atol = 1e-4 )
1655+
1656+ # Verify
1657+ for keepdims in [True , False ]:
1658+ # no `axes` input && `noop_with_empty_axes` = 0 -> reduce over all dimensions.
1659+ verify_reduce_func (
1660+ func ,
1661+ np .random .randn (3 , 2 , 2 ).astype (np .float32 ),
1662+ axes = [],
1663+ keepdims = keepdims ,
1664+ noop_with_empty_axes = False ,
1665+ )
1666+
1667+ # no `axes` input && `noop_with_empty_axes` = 0 -> reduce over all dimensions.
1668+ verify_reduce_func (
1669+ func ,
1670+ np .random .randn (3 , 2 , 2 ).astype (np .float32 ),
1671+ axes = None ,
1672+ keepdims = keepdims ,
1673+ noop_with_empty_axes = False ,
1674+ )
1675+
1676+ # no `axes` input && `noop_with_empty_axes` = 1 -> return the input unchanged.
1677+ verify_reduce_func (
1678+ func ,
1679+ np .random .randn (4 , 3 ).astype (np .float32 ),
1680+ axes = [],
1681+ keepdims = keepdims ,
1682+ noop_with_empty_axes = True ,
1683+ )
1684+
1685+ # no `axes` input && `noop_with_empty_axes` = 1 -> return the input unchanged.
1686+ # (onnxruntime bug) Runtime error on the onnxruntime part
1687+ # verify_reduce_func(
1688+ # func,
1689+ # np.random.randn(4, 3).astype(np.float32),
1690+ # axes=None,
1691+ # keepdims=keepdims,
1692+ # noop_with_empty_axes=True,
1693+ # )
1694+
1695+ # `axes` provided -> reduce over specified axes.
1696+ verify_reduce_func (
1697+ func ,
1698+ np .random .randn (3 , 3 , 3 , 1 ).astype (np .float32 ),
1699+ axes = (1 , 2 ),
1700+ keepdims = keepdims ,
1701+ noop_with_empty_axes = True ,
1702+ )
1703+
1704+
15801705@pytest .mark .parametrize ("in_dtype" , [np .float32 , np .int32 ])
15811706@pytest .mark .parametrize ("axis" , [None , 0 , 1 , 2 ])
15821707@pytest .mark .parametrize ("keepdims" , [None , True , False ])
0 commit comments