87 lines
2.9 KiB
Python
87 lines
2.9 KiB
Python
# Copyright (C) 2018-2024 Intel Corporation
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import pytest
|
|
import torch
|
|
import random
|
|
|
|
from pytorch_layer_test_class import PytorchLayerTest
|
|
|
|
|
|
class aten_prod(torch.nn.Module):
|
|
def __init__(self, in_dtype):
|
|
super(aten_prod, self).__init__()
|
|
self.in_dtype = in_dtype
|
|
|
|
def forward(self, x):
|
|
return torch.prod(x.to(self.in_dtype))
|
|
|
|
|
|
class aten_prod_dtype(torch.nn.Module):
|
|
def __init__(self, dtype, in_dtype):
|
|
super(aten_prod_dtype, self).__init__()
|
|
self.dtype = dtype
|
|
self.in_dtype = in_dtype
|
|
|
|
def forward(self, x):
|
|
return torch.prod(x.to(self.in_dtype))
|
|
|
|
|
|
class aten_prod_dim(torch.nn.Module):
|
|
def __init__(self, dim, keepdims, in_dtype):
|
|
super(aten_prod_dim, self).__init__()
|
|
self.dim = dim
|
|
self.keepdims = keepdims
|
|
self.in_dtype = in_dtype
|
|
|
|
def forward(self, x):
|
|
return torch.prod(x.to(self.in_dtype), self.dim, self.keepdims)
|
|
|
|
|
|
class aten_prod_dim_dtype(torch.nn.Module):
|
|
def __init__(self, dim, keepdims, dtype, in_dtype):
|
|
super(aten_prod_dim_dtype, self).__init__()
|
|
self.dim = dim
|
|
self.keepdims = keepdims
|
|
self.dtype = dtype
|
|
self.in_dtype = in_dtype
|
|
|
|
def forward(self, x):
|
|
return torch.prod(x.to(self.in_dtype), self.dim, self.keepdims, dtype=self.dtype)
|
|
|
|
|
|
class TestProd(PytorchLayerTest):
|
|
def _prepare_input(self, input_shape=(2), dtype=torch.float32):
|
|
import numpy as np
|
|
return (torch.randn(*input_shape).to(dtype).numpy(),)
|
|
|
|
@pytest.mark.parametrize("shape", [(1,),
|
|
(2,),
|
|
(2, 3),
|
|
(3, 4, 5),
|
|
(1, 2, 3, 4),
|
|
(1, 2, 3, 4, 5)])
|
|
@pytest.mark.parametrize("dtype", [None, torch.int32])
|
|
@pytest.mark.parametrize("in_dtype", [torch.float32, torch.bool])
|
|
@pytest.mark.parametrize("has_dim,keepdims", [(False, None), (True, True), (True, False)])
|
|
@pytest.mark.nightly
|
|
@pytest.mark.precommit
|
|
def test_prod(self, ie_device, precision, ir_version, shape, dtype, in_dtype, has_dim, keepdims):
|
|
if dtype is not None:
|
|
if has_dim:
|
|
m = aten_prod_dim_dtype(random.randint(0, len(shape) - 1),
|
|
keepdims,
|
|
dtype,
|
|
in_dtype)
|
|
else:
|
|
m = aten_prod_dtype(dtype, in_dtype)
|
|
else:
|
|
if has_dim:
|
|
m = aten_prod_dim(random.randint(0, len(shape) - 1),
|
|
keepdims,
|
|
in_dtype)
|
|
else:
|
|
m = aten_prod(in_dtype)
|
|
self._test(m, None, 'aten::prod', ie_device, precision, ir_version,
|
|
kwargs_to_prepare_input={'input_shape': shape, 'dtype': in_dtype})
|