diff --git a/smartlearner/batch_scheduler.py b/smartlearner/batch_scheduler.py
index 664cfb5..99f33ec 100644
--- a/smartlearner/batch_scheduler.py
+++ b/smartlearner/batch_scheduler.py
@@ -26,6 +26,11 @@ class MiniBatchScheduler(BatchScheduler):
         self.batch_size = batch_size
         self.shared_batch_count = theano.shared(np.array(0, dtype='i4'))
 
+        # Keep only `batch_size` examples as test values.
+        self.dataset.symb_inputs.tag.test_value = self.dataset.inputs.get_value()[:batch_size]
+        if self.dataset.has_targets:
+            self.dataset.symb_targets.tag.test_value = self.dataset.targets.get_value()[:batch_size]
+
     @property
     def batch_size(self):
         return self._shared_batch_size.get_value()
diff --git a/smartlearner/interfaces/dataset.py b/smartlearner/interfaces/dataset.py
index d2d0a4c..c67da77 100644
--- a/smartlearner/interfaces/dataset.py
+++ b/smartlearner/interfaces/dataset.py
@@ -1,13 +1,48 @@
-import theano
+import numpy as np
+
+import theano.tensor as T
+
+from smartlearner.utils import sharedX
 
 
 class Dataset(object):
+    """ Dataset interface.
+
+    Attributes
+    ----------
+    symb_inputs : `theano.tensor.TensorType` object
+        Symbolic variables representing the inputs.
+    symb_targets : `theano.tensor.TensorType` object or None
+        Symbolic variables representing the targets.
+
+    Notes
+    -----
+    `symb_inputs` and `symb_targets` have test value already tagged to them. Use
+    THEANO_FLAGS="compute_test_value=warn" to use them.
+    """
     def __init__(self, inputs, targets=None, name="dataset"):
+        """
+        Parameters
+        ----------
+        inputs : ndarray
+            Training examples
+        targets : ndarray (optional)
+            Target for each training example.
+        name : str (optional)
+            The name of the dataset is used to name Theano variables. Default: 'dataset'.
+        """
         self.name = name
         self.inputs = inputs
         self.targets = targets
-        self.symb_inputs = theano.tensor.matrix(name=self.name+'_inputs')
-        self.symb_targets = theano.tensor.matrix(name=self.name+'_targets')
+        self.symb_inputs = T.TensorVariable(type=T.TensorType("floatX", [False]*self.inputs.ndim),
+                                            name=self.name+'_symb_inputs')
+        self.symb_inputs.tag.test_value = self.inputs.get_value()  # For debugging Theano graphs.
+
+        self.symb_targets = None
+        if self.has_targets:
+            self.symb_targets = T.TensorVariable(type=T.TensorType("floatX", [False]*self.targets.ndim),
+                                                 name=self.name+'_symb_targets')
+            self.symb_targets.tag.test_value = self.targets.get_value()  # For debugging Theano graphs.
 
     @property
     def inputs(self):
@@ -15,7 +50,7 @@ class Dataset(object):
 
     @inputs.setter
     def inputs(self, value):
-        self._inputs_shared = theano.shared(value, name=self.name + "_inputs", borrow=True)
+        self._inputs_shared = sharedX(value, name=self.name+"_inputs")
 
     @property
     def targets(self):
@@ -24,20 +59,37 @@ class Dataset(object):
     @targets.setter
     def targets(self, value):
         if value is not None:
-            self._targets_shared = theano.shared(value, name=self.name + "_targets", borrow=True)
+            self._targets_shared = sharedX(np.array(value), name=self.name+"_targets")
         else:
             self._targets_shared = None
 
+    @property
+    def has_targets(self):
+        return self.targets is not None
+
+    @property
+    def input_shape(self):
+        return self.inputs.get_value().shape[1:]
+
+    @property
+    def target_shape(self):
+        if self.has_targets:
+            return self.targets.get_value().shape[1:]
+
+        return None
+
     @property
     def input_size(self):
-        return self.inputs.get_value().shape[-1]
+        # TODO: is this property really useful? If needed one could just call directly `dataset.input_shape[-1]`.
+        return self.input_shape[-1]
 
     @property
     def target_size(self):
-        if self.targets is None:
-            return 0
-        else:
-            return self.targets.get_value().shape[-1]
+        # TODO: is this property really useful? If needed one could just call directly `dataset.target_shape[-1]`.
+        if self.has_targets:
+            return self.target_shape[-1]
+
+        return None
 
     def __len__(self):
         return len(self.inputs.get_value())

Move tests at the root of the library.
That way test won't appear in the autocomplete of your preferred python IDE.
diff --git a/tests/interfaces/test_dataset.py b/tests/interfaces/test_dataset.py
new file mode 100644
index 0000000..6378b93
--- /dev/null
+++ b/tests/interfaces/test_dataset.py
@@ -0,0 +1,140 @@
+import numpy as np
+import theano
+import theano.tensor as T
+
+from nose.tools import assert_true
+from numpy.testing import assert_equal, assert_array_equal
+
+from smartlearner.interfaces.dataset import Dataset
+
+floatX = theano.config.floatX
+ALL_DTYPES = np.sctypes['int'] + np.sctypes['uint'] + np.sctypes['float']
+
+
+def test_dataset_used_in_theano_function():
+    rng = np.random.RandomState(1234)
+
+    nb_examples = 10
+
+    inputs = (rng.randn(nb_examples, 5) * 100).astype(floatX)
+    targets = (rng.randn(nb_examples, 1) > 0.5).astype(floatX)
+    dataset = Dataset(inputs, targets)
+
+    input_sqr_norm = T.sum(dataset.symb_inputs**2)
+    result = input_sqr_norm - dataset.symb_targets
+    f = theano.function([dataset.symb_inputs, dataset.symb_targets], result)
+
+    assert_array_equal(f(inputs, targets), np.sum(inputs**2)-targets)
+
+
+def test_dataset_without_targets():
+    rng = np.random.RandomState(1234)
+
+    nb_examples = 10
+    nb_features = 3
+    sequences_length = 4
+    nb_channels = 2
+    image_shape = (5, 5)
+
+    # Test creating dataset with different example shapes:
+    # scalar feature, vector features, sequence of vector features, multiple channels images features.
+    for example_shape in [(), (nb_features,), (sequences_length, nb_features), (nb_channels,)+image_shape]:
+        inputs_shape = (nb_examples,) + example_shape
+
+        for dtype in ALL_DTYPES:
+            inputs = (rng.randn(*inputs_shape) * 100).astype(dtype)
+            dataset = Dataset(inputs)
+
+            # Data should be converted into `floatX`.
+            assert_equal(dataset.inputs.dtype, floatX)
+            assert_equal(dataset.symb_inputs.dtype, floatX)
+            assert_equal(dataset.symb_inputs.ndim, inputs.ndim)
+            assert_equal(dataset.input_shape, example_shape)
+            assert_array_equal(dataset.inputs.get_value(), inputs.astype(floatX))
+
+            # Everything related to target should be None
+            assert_true(dataset.targets is None)
+            assert_true(dataset.symb_targets is None)
+            assert_true(dataset.target_shape is None)
+            assert_true(dataset.target_size is None)
+
+    # Create dataset from nested Pyton lists.
+    inputs = [[1, 2, 3]] * nb_examples
+    dataset = Dataset(inputs)
+    # Data should be converted into `floatX`.
+    assert_equal(dataset.inputs.dtype, floatX)
+    assert_equal(dataset.symb_inputs.dtype, floatX)
+    assert_equal(dataset.symb_inputs.ndim, 2)
+    assert_equal(dataset.input_shape, (3,))
+    assert_array_equal(dataset.inputs.get_value(), np.array(inputs, dtype=floatX))
+
+
+def test_dataset_with_targets():
+    rng = np.random.RandomState(1234)
+
+    nb_examples = 10
+    nb_features = 3
+    sequences_length = 4
+    nb_channels = 2
+    image_shape = (5, 5)
+
+    # Test creating dataset with different example shapes and target shapes:
+    # scalar feature, vector features, sequence of vector features, multiple channels images features.
+    for target_shape in [(), (nb_features,), (sequences_length, nb_features), (nb_channels,)+image_shape]:
+        for example_shape in [(), (nb_features,), (sequences_length, nb_features), (nb_channels,)+image_shape]:
+            inputs_shape = (nb_examples,) + example_shape
+            targets_shape = (nb_examples,) + target_shape
+
+            for example_dtype in ALL_DTYPES:
+                for target_dtype in ALL_DTYPES:
+                    inputs = (rng.randn(*inputs_shape) * 100).astype(example_dtype)
+                    targets = (rng.randn(*targets_shape) * 100).astype(target_dtype)
+                    dataset = Dataset(inputs, targets)
+
+                    # Data should be converted into `floatX`.
+                    assert_equal(dataset.inputs.dtype, floatX)
+                    assert_equal(dataset.symb_inputs.dtype, floatX)
+                    assert_equal(dataset.symb_inputs.ndim, inputs.ndim)
+                    assert_equal(dataset.input_shape, example_shape)
+                    assert_array_equal(dataset.inputs.get_value(), inputs.astype(floatX))
+
+                    assert_equal(dataset.targets.dtype, floatX)
+                    assert_equal(dataset.symb_targets.dtype, floatX)
+                    assert_equal(dataset.symb_targets.ndim, targets.ndim)
+                    assert_equal(dataset.target_shape, target_shape)
+                    assert_array_equal(dataset.targets.get_value(), targets.astype(floatX))
+
+    # Create dataset from nested Pyton lists.
+    inputs = [[1, 2, 3]] * nb_examples
+    targets = [[1, 2, 3]] * nb_examples
+    dataset = Dataset(inputs, targets)
+    # Data should be converted into `floatX`.
+    assert_equal(dataset.inputs.dtype, floatX)
+    assert_equal(dataset.symb_inputs.dtype, floatX)
+    assert_equal(dataset.symb_inputs.ndim, 2)
+    assert_equal(dataset.input_shape, (3,))
+    assert_array_equal(dataset.inputs.get_value(), np.array(inputs, dtype=floatX))
+
+    assert_equal(dataset.targets.dtype, floatX)
+    assert_equal(dataset.symb_targets.dtype, floatX)
+    assert_equal(dataset.symb_targets.ndim, 2)
+    assert_equal(dataset.target_shape, (3,))
+    assert_array_equal(dataset.targets.get_value(), np.array(targets, dtype=floatX))
+
+
+def test_dataset_with_test_value():
+    rng = np.random.RandomState(1234)
+
+    nb_examples = 10
+
+    theano.config.compute_test_value = 'warn'
+    try:
+        inputs = (rng.randn(nb_examples, 5) * 100).astype(floatX)
+        targets = (rng.randn(nb_examples, 1) > 0.5).astype(floatX)
+        dataset = Dataset(inputs, targets)
+
+        input_sqr_norm = T.sum(dataset.symb_inputs**2)
+        result = input_sqr_norm - dataset.symb_targets
+        assert_array_equal(result.tag.test_value, np.sum(inputs**2)-targets)
+    finally:
+        theano.config.compute_test_value = 'off'

pip install -e .
parse_log_pytest
theano numpy scipy nose pyparsing pip flake8 six pep8 pyflakes
nose nose-cov pytest
3.9
pytest --no-header -rA --tb=line --color=no -p no:cacheprovider -W ignore::DeprecationWarning
cov-core==1.15.0
coverage==7.8.0
exceptiongroup==1.2.2
flake8 @ file:///croot/flake8_1726157165993/work
iniconfig==2.1.0
Mako @ file:///croot/mako_1665472421453/work
MarkupSafe @ file:///croot/markupsafe_1738584038848/work
mccabe @ file:///opt/conda/conda-bld/mccabe_1644221741721/work
nose @ file:///opt/conda/conda-bld/nose_1642704612149/work
nose-cov==1.6
numpy @ file:///opt/conda/conda-bld/numpy_and_numpy_base_1652801679809/work
packaging==24.2
pep8==1.7.1
pluggy==1.5.0
pycodestyle @ file:///croot/pycodestyle_1726150303809/work
pyflakes @ file:///croot/pyflakes_1708962956225/work
pygpu==0.7.6
pyparsing @ file:///croot/pyparsing_1731445506121/work
pytest==8.3.5
scipy @ file:///opt/conda/conda-bld/scipy_1661390393401/work
six @ file:///tmp/build/80754af9/six_1644875935023/work
-e git+https://github.com/SMART-Lab/smartlearner.git@b0877d3b961deceb273139f064985e39239351a7#egg=smartlearner
Theano==1.0.5
tomli==2.2.1

name: smartlearner
channels:
  - defaults
  - https://repo.anaconda.com/pkgs/main
  - https://repo.anaconda.com/pkgs/r
  - conda-forge
dependencies:
  - _libgcc_mutex=0.1=main
  - _openmp_mutex=5.1=1_gnu
  - _sysroot_linux-64_curr_repodata_hack=3=haa98f57_10
  - binutils_impl_linux-64=2.35.1=h27ae35d_9
  - binutils_linux-64=2.35.1=h454624a_30
  - blas=1.0=openblas
  - ca-certificates=2025.2.25=h06a4308_0
  - fftw=3.3.9=h5eee18b_2
  - flake8=7.1.1=py39h06a4308_0
  - gcc_impl_linux-64=7.5.0=h7105cf2_17
  - gcc_linux-64=7.5.0=h8f34230_30
  - gxx_impl_linux-64=7.5.0=h0a5bf11_17
  - gxx_linux-64=7.5.0=hffc177d_30
  - kernel-headers_linux-64=3.10.0=h57e8cba_10
  - ld_impl_linux-64=2.35.1=h7274673_9
  - libffi=3.4.4=h6a678d5_1
  - libgcc-devel_linux-64=7.5.0=hbbeae57_17
  - libgcc-ng=11.2.0=h1234567_1
  - libgfortran-ng=11.2.0=h00389a5_1
  - libgfortran5=11.2.0=h1234567_1
  - libgomp=11.2.0=h1234567_1
  - libgpuarray=0.7.6=h7f8727e_1
  - libopenblas=0.3.21=h043d6bf_0
  - libstdcxx-devel_linux-64=7.5.0=hf0c5c8d_17
  - libstdcxx-ng=11.2.0=h1234567_1
  - mako=1.2.3=py39h06a4308_0
  - markupsafe=3.0.2=py39h5eee18b_0
  - mccabe=0.7.0=pyhd3eb1b0_0
  - ncurses=6.4=h6a678d5_0
  - nose=1.3.7=pyhd3eb1b0_1008
  - numpy=1.22.3=py39h7a5d4dd_0
  - numpy-base=1.22.3=py39hb8be1f0_0
  - openssl=3.0.16=h5eee18b_0
  - pep8=1.7.1=py39h06a4308_1
  - pip=25.0=py39h06a4308_0
  - pycodestyle=2.12.1=py39h06a4308_0
  - pyflakes=3.2.0=py39h06a4308_0
  - pygpu=0.7.6=py39hce1f21e_1
  - pyparsing=3.2.0=py39h06a4308_0
  - python=3.9.21=he870216_1
  - readline=8.2=h5eee18b_0
  - scipy=1.7.3=py39hf838250_2
  - setuptools=75.8.0=py39h06a4308_0
  - six=1.16.0=pyhd3eb1b0_1
  - sqlite=3.45.3=h5eee18b_0
  - sysroot_linux-64=2.17=h57e8cba_10
  - theano=1.0.5=py39h295c915_1
  - tk=8.6.14=h39e8969_0
  - tzdata=2025a=h04d1e81_0
  - wheel=0.45.1=py39h06a4308_0
  - xz=5.6.4=h5eee18b_1
  - zlib=1.2.13=h5eee18b_1
  - pip:
      - cov-core==1.15.0
      - coverage==7.8.0
      - exceptiongroup==1.2.2
      - iniconfig==2.1.0
      - nose-cov==1.6
      - packaging==24.2
      - pluggy==1.5.0
      - pytest==8.3.5
      - tomli==2.2.1
prefix: /opt/conda/envs/smartlearner

tests/interfaces/test_dataset.py::test_dataset_without_targets tests/interfaces/test_dataset.py::test_dataset_with_targets
tests/interfaces/test_dataset.py::test_dataset_used_in_theano_function tests/interfaces/test_dataset.py::test_dataset_with_test_value
smartlearner/interfaces/dataset.py smartlearner/batch_scheduler.py
smartlearner/interfaces/dataset.py smartlearner/batch_scheduler.py