<filename>tests/model_tests.py
import sys
import unittest
import tensorflow as tf

import tfdiffeq
import tfdiffeq.models as models

if not tf.executing_eagerly():
    tf.enable_eager_execution()


class TestModel(unittest.TestCase):

    def test_dense_odenet(self):
        model = models.ODENet(hidden_dim=10, output_dim=2)

        data = tf.zeros([10, 5])
        out = model(data)

        assert out.shape == (10, 2)
        assert len(model.layers) == 2

    def test_dense_odenet_time_dependent(self):
        model = models.ODENet(hidden_dim=10, output_dim=2, time_dependent=True)

        data = tf.zeros([10, 5])
        out = model(data)

        assert out.shape == (10, 2)
        assert len(model.layers) == 2

    def test_conv_odenet(self):
        model = models.Conv2dODENet(num_filters=10, output_dim=2)

        data = tf.zeros([10, 5, 5, 3])
        out = model(data)

        assert out.shape == (10, 5, 5, 2)
        assert len(model.layers) == 2

    def test_conv_odenet_time_dependent(self):
        model = models.Conv2dODENet(num_filters=10, output_dim=2, time_dependent=True)

        data = tf.zeros([10, 5, 5, 3])
        out = model(data)

        assert out.shape == (10, 5, 5, 2)
        assert len(model.layers) == 2


if __name__ == '__main__':
    unittest.main(__file__)

<filename>tests/model_tests.py
import sys
import unittest
import tensorflow as tf

import tfdiffeq
import tfdiffeq.models as models

if not tf.executing_eagerly():
    tf.enable_eager_execution()


class TestModel(unittest.TestCase):

    def test_dense_odenet(self):
        model = models.ODENet(hidden_dim=10, output_dim=2)

        data = tf.zeros([10, 5])
        out = model(data)

        assert out.shape == (10, 2)
        assert len(model.layers) == 2

    def test_dense_odenet_time_dependent(self):
        model = models.ODENet(hidden_dim=10, output_dim=2, time_dependent=True)

        data = tf.zeros([10, 5])
        out = model(data)

        assert out.shape == (10, 2)
        assert len(model.layers) == 2

    def test_conv_odenet(self):
        model = models.Conv2dODENet(num_filters=10, output_dim=2)

        data = tf.zeros([10, 5, 5, 3])
        out = model(data)

        assert out.shape == (10, 5, 5, 2)
        assert len(model.layers) == 2

    def test_conv_odenet_time_dependent(self):
        model = models.Conv2dODENet(num_filters=10, output_dim=2, time_dependent=True)

        data = tf.zeros([10, 5, 5, 3])
        out = model(data)

        assert out.shape == (10, 5, 5, 2)
        assert len(model.layers) == 2


if __name__ == '__main__':
    unittest.main(__file__)
