from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import torch
import unittest
from gpytorch.means import ZeroMean


class TestZeroMean(unittest.TestCase):
    def setUp(self):
        self.mean = ZeroMean()

    def test_forward(self):
        a = torch.Tensor([[1, 2], [2, 4]])
        res = self.mean(a)
        self.assertEqual(tuple(res.size()), (2,))
        self.assertTrue(res.eq(0).all())

    def test_forward_batch(self):
        a = torch.Tensor([[[1, 2], [1, 2], [2, 4]], [[2, 3], [2, 3], [1, 3]]])
        res = self.mean(a)
        self.assertEqual(tuple(res.size()), (2, 3))
        self.assertTrue(res.eq(0).all())

[((13, 20, 13, 30), 'gpytorch.means.ZeroMean', 'ZeroMean', ({}, {}), '()', False, 'from gpytorch.means import ZeroMean\n'), ((16, 12, 16, 42), 'torch.Tensor', 'torch.Tensor', ({(16, 25, 16, 41): '[[1, 2], [2, 4]]'}, {}), '([[1, 2], [2, 4]])', False, 'import torch\n'), ((22, 12, 22, 78), 'torch.Tensor', 'torch.Tensor', ({(22, 25, 22, 77): '[[[1, 2], [1, 2], [2, 4]], [[2, 3], [2, 3], [1, 3]]]'}, {}), '([[[1, 2], [1, 2], [2, 4]], [[2, 3], [2, 3], [1, 3]]])', False, 'import torch\n')]