import unittest
import torch
import importlib
import io
import sys


def _dynamic(test, cond, ok, fail):
    if cond:
        test._testMethodName = ok
        test.assertTrue(True)
    else:
        test._testMethodName = fail
        test.fail(fail)


class TestUserCode(unittest.TestCase):

    def setUp(self):
        # Reference tensor for comparison
        self.tensor = torch.tensor([[3, 7, 1],
                                    [2, 5, 9],
                                    [8, 6, 4]], dtype=torch.float32)

    # --- Variable existence ---
    def test_row_sum_exists(self):
        import user_code
        cond = hasattr(user_code, "row_sum")
        _dynamic(
            self,
            cond,
            "`row_sum` variable is defined.",
            "`row_sum` is missing. Compute it using `tensor.sum(dim=1, keepdim=True)`."
        )

    def test_column_mean_exists(self):
        import user_code
        cond = hasattr(user_code, "column_mean")
        _dynamic(
            self,
            cond,
            "`column_mean` variable is defined.",
            "`column_mean` is missing. Compute it using `tensor.mean(dim=0)`."
        )

    def test_row_max_exists(self):
        import user_code
        cond = hasattr(user_code, "row_max")
        _dynamic(
            self,
            cond,
            "`row_max` variable is defined.",
            "`row_max` is missing. Compute it using `tensor.max(dim=1)[0]`."
        )

    def test_normalized_exists(self):
        import user_code
        cond = hasattr(user_code, "normalized_tensor")
        _dynamic(
            self,
            cond,
            "`normalized_tensor` variable is defined.",
            "`normalized_tensor` is missing. Normalize using `tensor / tensor.sum()`."
        )

    # --- Correct calculations ---
    def test_row_sum_correct(self):
        import user_code
        expected = self.tensor.sum(dim=1, keepdim=True)
        cond = torch.allclose(user_code.row_sum, expected)
        _dynamic(
            self,
            cond,
            "`row_sum` is computed correctly.",
            "`row_sum` is incorrect. Use `tensor.sum(dim=1, keepdim=True)`."
        )

    def test_column_mean_correct(self):
        import user_code
        expected = self.tensor.mean(dim=0)
        cond = torch.allclose(user_code.column_mean, expected)
        _dynamic(
            self,
            cond,
            "`column_mean` is computed correctly.",
            "`column_mean` is incorrect. Use `tensor.mean(dim=0)`."
        )

    def test_row_max_correct(self):
        import user_code
        expected = self.tensor.max(dim=1)[0]
        cond = torch.allclose(user_code.row_max, expected)
        _dynamic(
            self,
            cond,
            "`row_max` is computed correctly.",
            "`row_max` is incorrect. Use `tensor.max(dim=1)[0]`."
        )

    def test_normalized_tensor_correct(self):
        import user_code
        expected = self.tensor / self.tensor.sum()
        cond = torch.allclose(user_code.normalized_tensor, expected)
        _dynamic(
            self,
            cond,
            "`normalized_tensor` is computed correctly.",
            "`normalized_tensor` is incorrect. Normalize using `tensor / tensor.sum()`."
        )

    # --- Ensure something was printed ---
    def test_print_output(self):
        import user_code

        buffer = io.StringIO()
        sys.stdout = buffer
        importlib.reload(user_code)
        sys.stdout = sys.__stdout__

        cond = len(buffer.getvalue().strip()) > 0

        _dynamic(
            self,
            cond,
            "Output is successfully printed.",
            "No output detected. Make sure to print all required values."
        )


if __name__ == "__main__":
    unittest.main()


test_code.py

Learn the fundamental and advanced concepts needed to work with PyTorch efficiently. Gain a solid understanding tensors, including creation, operations, and reshaping. Explore the essentials of gradients, backpropagation, and linear regression before moving on to handling datasets. Master the skill needed build, train, and evaluate neural networks.

Explore the fundamentals of PyTorch, focusing on tensors—the core data structure used for computations. You will learn about tensor creation, random initialization, mathematical operations, and shape manipulation.

Explore key concepts for training models in PyTorch, including computing gradients and performing multi-step backpropagation. Master linear regression as a foundational machine learning model and introduces handling datasets efficiently.

Discover how to build, train, and evaluate neural networks using PyTorch. You will learn how to define a simple feedforward neural network, optimize its parameters through training, and assess its performance.

Challenge: Performing Mathematical Operations

解答