Repository: nipype/pydra

Issue:
State changes: adding container_dim concept 
We want to add an optional argument to the `split` that specify the dimension of the "container" that will be used for splitting.  This is important to be able to recognize what is the dimension of container and what is the dimension of the object that should be passed to the function when input has multiple dimensions, e.g. `a = [[1, 2], [10, 20]]`. We can have a function that takes a single number, i.e. `1`, `2`, `10`, and `20`  or a list, i.e., `[1, 2]` and `[10, 20]`. Right now the second case  has to be represented by using tuples, i.e. `a=[(1, 2), (10, 20)]` (see [an example](https://github.com/nipype/pydra/blob/master/pydra/engine/tests/test_node_task.py#L646)).

We want to assume that if the input is a list (possibly also a tuple), the default `container_dim` is `1`, so we iterate `for el in a:`. However, the user can specify `container_dim=2` to iterate `for el_l in a: for el in el_l:`. Higher container dimension will be also used for the output of the task that have outer splitter.

We also want to allow to iterate over string - as a default a single string should have `container_dim=0`, but if the user specifies `conatiner_dim=1`, pydra will iterate over the letters. 

@satra - please comment if I forgot something important

**edit**
`container_dim` should be per input, so probably a dictionary

Discussion:

Patch:
diff --git a/pydra/engine/helpers_state.py b/pydra/engine/helpers_state.py
index 65a580de..8136e2a1 100644
--- a/pydra/engine/helpers_state.py
+++ b/pydra/engine/helpers_state.py
@@ -362,14 +362,15 @@ def iter_splits(iterable, keys):
         yield dict(zip(keys, list(flatten(iter, max_depth=1000))))
 
 
-def input_shape(in1):
-    """Get input shape."""
+def input_shape(inp, cont_dim=1):
+    """Get input shape, depends on the container dimension, if not specify it is assumed to be 1 """
     # TODO: have to be changed for inner splitter (sometimes different length)
-    shape = [len(in1)]
+    cont_dim -= 1
+    shape = [len(inp)]
     last_shape = None
-    for value in in1:
-        if isinstance(value, list):
-            cur_shape = input_shape(value)
+    for value in inp:
+        if isinstance(value, list) and cont_dim > 0:
+            cur_shape = input_shape(value, cont_dim)
             if last_shape is None:
                 last_shape = cur_shape
             elif last_shape != cur_shape:
@@ -383,11 +384,37 @@ def input_shape(in1):
     return tuple(shape)
 
 
-def splits(splitter_rpn, inputs, inner_inputs=None):
-    """Split process as specified by an rpn splitter, from left to right."""
+def splits(splitter_rpn, inputs, inner_inputs=None, cont_dim=None):
+    """
+    Splits input variable as specified by splitter
+
+    Parameters
+    ----------
+    splitter_rpn : list
+        splitter in RPN notation
+    inputs: dict
+        input variables
+    inner_inputs: dict, optional
+        inner input specification
+    cont_dim: dict, optional
+        container dimension for input variable, specifies how nested is the intput,
+        if not specified 1 will be used for all inputs (so will not be flatten)
+
+
+    Returns
+    -------
+    splitter : list
+        each element contains indices for inputs
+    keys: list
+        names of input variables
+
+    """
+
     stack = []
     keys = []
-    shapes_var = {}
+    if cont_dim is None:
+        cont_dim = {}
+    # analysing states from connected tasks if inner_inputs
     if inner_inputs:
         previous_states_ind = {
             "_{}".format(v.name): (v.ind_l_final, v.keys_final)
@@ -407,9 +434,9 @@ def splits(splitter_rpn, inputs, inner_inputs=None):
             op_single,
             inputs,
             inner_inputs,
-            shapes_var,
             previous_states_ind,
             keys_fromLeftSpl,
+            cont_dim=cont_dim,
         )
 
     terms = {}
@@ -418,7 +445,11 @@ def splits(splitter_rpn, inputs, inner_inputs=None):
     shape = {}
     # iterating splitter_rpn
     for token in splitter_rpn:
-        if token in [".", "*"]:
+        if token not in [".", "*"]:  # token is one of the input var
+            # adding variable to the stack
+            stack.append(token)
+        else:
+            # removing Right and Left var from the stack
             terms["R"] = stack.pop()
             terms["L"] = stack.pop()
             # checking if terms are strings, shapes, etc.
@@ -429,10 +460,14 @@ def splits(splitter_rpn, inputs, inner_inputs=None):
                         trm_val[lr] = previous_states_ind[term][0]
                         shape[lr] = (len(trm_val[lr]),)
                     else:
-                        shape[lr] = input_shape(inputs[term])
+                        if term in cont_dim:
+                            shape[lr] = input_shape(
+                                inputs[term], cont_dim=cont_dim[term]
+                            )
+                        else:
+                            shape[lr] = input_shape(inputs[term])
                         trm_val[lr] = range(reduce(lambda x, y: x * y, shape[lr]))
                     trm_str[lr] = True
-                    shapes_var[term] = shape[lr]
                 else:
                     trm_val[lr], shape[lr] = term
                     trm_str[lr] = False
@@ -447,6 +482,7 @@ def splits(splitter_rpn, inputs, inner_inputs=None):
                     )
                 newshape = shape["R"]
             if token == "*":
+                # TODO: pomyslec
                 newshape = tuple(list(shape["L"]) + list(shape["R"]))
 
             # creating list with keys
@@ -466,7 +502,6 @@ def splits(splitter_rpn, inputs, inner_inputs=None):
             elif trm_str["R"]:
                 keys = keys + new_keys["R"]
 
-            #
             newtrm_val = {}
             for lr in ["R", "L"]:
                 # TODO: rewrite once I have more tests
@@ -491,13 +526,11 @@ def splits(splitter_rpn, inputs, inner_inputs=None):
 
             pushval = (op[token](newtrm_val["L"], newtrm_val["R"]), newshape)
             stack.append(pushval)
-        else:  # name of one of the inputs (token not in [".", "*"])
-            stack.append(token)
 
     val = stack.pop()
     if isinstance(val, tuple):
         val = val[0]
-    return val, keys, shapes_var, keys_fromLeftSpl
+    return val, keys, keys_fromLeftSpl
 
 
 # dj: TODO: do I need keys?
@@ -636,17 +669,22 @@ def splits_groups(splitter_rpn, combiner=None, inner_inputs=None):
 
 
 def _single_op_splits(
-    op_single, inputs, inner_inputs, shapes_var, previous_states_ind, keys_fromLeftSpl
+    op_single,
+    inputs,
+    inner_inputs,
+    previous_states_ind,
+    keys_fromLeftSpl,
+    cont_dim=None,
 ):
     if op_single.startswith("_"):
         return (
             previous_states_ind[op_single][0],
             previous_states_ind[op_single][1],
-            None,
             keys_fromLeftSpl,
         )
-    shape = input_shape(inputs[op_single])
-    shapes_var[op_single] = shape
+    if cont_dim is None:
+        cont_dim = {}
+    shape = input_shape(inputs[op_single], cont_dim=cont_dim.get(op_single, 1))
     trmval = range(reduce(lambda x, y: x * y, shape))
     if op_single in inner_inputs:
         # TODO: have to be changed if differ length
@@ -659,11 +697,11 @@ def _single_op_splits(
         res = op["."](op_out, trmval)
         val = res
         keys = inner_inputs[op_single].keys_final + [op_single]
-        return val, keys, shapes_var, keys_fromLeftSpl
+        return val, keys, keys_fromLeftSpl
     else:
         val = op["*"](trmval)
         keys = [op_single]
-        return val, keys, shapes_var, keys_fromLeftSpl
+        return val, keys, keys_fromLeftSpl
 
 
 def _single_op_splits_groups(
@@ -727,10 +765,15 @@ def combine_final_groups(combiner, groups, groups_stack, keys):
     return keys_final, groups_final, groups_stack_final, combiner_all
 
 
-def map_splits(split_iter, inputs):
+def map_splits(split_iter, inputs, cont_dim=None):
     """Get a dictionary of prescribed splits."""
+    if cont_dim is None:
+        cont_dim = {}
     for split in split_iter:
-        yield {k: list(flatten(ensure_list(inputs[k])))[v] for k, v in split.items()}
+        yield {
+            k: list(flatten(ensure_list(inputs[k]), max_depth=cont_dim.get(k, None)))[v]
+            for k, v in split.items()
+        }
 
 
 # Functions for merging and completing splitters in states.
diff --git a/pydra/engine/state.py b/pydra/engine/state.py
index 4d19abc0..e41eb7cf 100644
--- a/pydra/engine/state.py
+++ b/pydra/engine/state.py
@@ -346,7 +346,7 @@ class State:
                 stack = [gr + nr_gr_f for gr in stack]
                 self.groups_stack_final.append(stack)
 
-    def prepare_states(self, inputs):
+    def prepare_states(self, inputs, cont_dim=None):
         """
         Prepare a full list of state indices and state values.
 
@@ -357,6 +357,11 @@ class State:
             specific elements from inputs that can be used running interfaces
 
         """
+        # container dimension for each input, specifies how nested the input is
+        if cont_dim is None:
+            self.cont_dim = {}
+        else:
+            self.cont_dim = cont_dim
         if isinstance(inputs, BaseSpec):
             self.inputs = hlpst.inputs_types_to_dict(self.name, inputs)
         else:
@@ -366,7 +371,7 @@ class State:
                 # I think now this if is never used
                 if not hasattr(st, "states_ind"):
                     # dj: should i provide different inputs?
-                    st.prepare_states(self.inputs)
+                    st.prepare_states(self.inputs, cont_dim=cont_dim)
                 self.inputs.update(st.inputs)
         self.prepare_states_ind()
         self.prepare_states_val()
@@ -395,8 +400,11 @@ class State:
         partial_rpn = hlpst.remove_inp_from_splitter_rpn(
             deepcopy(self.splitter_rpn_compact), elements_to_remove
         )
-        values_out_pr, keys_out_pr, _, kL = hlpst.splits(
-            partial_rpn, self.inputs, inner_inputs=self.inner_inputs
+        values_out_pr, keys_out_pr, kL = hlpst.splits(
+            partial_rpn,
+            self.inputs,
+            inner_inputs=self.inner_inputs,
+            cont_dim=self.cont_dim,
         )
         values_pr = list(values_out_pr)
 
@@ -429,8 +437,11 @@ class State:
         )
         # TODO: create a function for this!!
         if combined_rpn:
-            val_r, key_r, _, _ = hlpst.splits(
-                combined_rpn, self.inputs, inner_inputs=self.inner_inputs
+            val_r, key_r, _ = hlpst.splits(
+                combined_rpn,
+                self.inputs,
+                inner_inputs=self.inner_inputs,
+                cont_dim=self.cont_dim,
             )
             values = list(val_r)
         else:
@@ -464,7 +475,9 @@ class State:
 
     def prepare_states_val(self):
         """Evaluate states values having states indices."""
-        self.states_val = list(hlpst.map_splits(self.states_ind, self.inputs))
+        self.states_val = list(
+            hlpst.map_splits(self.states_ind, self.inputs, cont_dim=self.cont_dim)
+        )
         return self.states_val
 
     def prepare_inputs(self):
@@ -489,8 +502,11 @@ class State:
                 deepcopy(self.splitter_rpn_compact), elements_to_remove
             )
             if partial_rpn:
-                values_inp, keys_inp, _, _ = hlpst.splits(
-                    partial_rpn, self.inputs, inner_inputs=self.inner_inputs
+                values_inp, keys_inp, _ = hlpst.splits(
+                    partial_rpn,
+                    self.inputs,
+                    inner_inputs=self.inner_inputs,
+                    cont_dim=self.cont_dim,
                 )
                 inputs_ind = values_inp
             else:


Test patch:
diff --git a/pydra/engine/tests/test_helpers_state.py b/pydra/engine/tests/test_helpers_state.py
index ce796cf7..f1eae8d9 100644
--- a/pydra/engine/tests/test_helpers_state.py
+++ b/pydra/engine/tests/test_helpers_state.py
@@ -91,17 +91,19 @@ def test_splits_groups_comb(
 
 
 @pytest.mark.parametrize(
-    "splitter, values, keys, splits",
+    "splitter, cont_dim, values, keys, splits",
     [
-        ("a", [(0,), (1,)], ["a"], [{"a": 1}, {"a": 2}]),
+        ("a", None, [(0,), (1,)], ["a"], [{"a": 1}, {"a": 2}]),
         (
             ("a", "v"),
+            None,
             [(0, 0), (1, 1)],
             ["a", "v"],
             [{"a": 1, "v": "a"}, {"a": 2, "v": "b"}],
         ),
         (
             ["a", "v"],
+            None,
             [(0, 0), (0, 1), (1, 0), (1, 1)],
             ["a", "v"],
             [
@@ -113,24 +115,28 @@ def test_splits_groups_comb(
         ),
         (
             ("a", "v", "c"),
+            None,
             [((0, 0), 0), ((1, 1), 1)],
             ["a", "v", "c"],
             [{"a": 1, "c": 3, "v": "a"}, {"a": 2, "c": 4, "v": "b"}],
         ),
         (
             (("a", "v"), "c"),
+            None,
             [((0, 0), 0), ((1, 1), 1)],
             ["a", "v", "c"],
             [{"a": 1, "c": 3, "v": "a"}, {"a": 2, "c": 4, "v": "b"}],
         ),
         (
             ("a", ("v", "c")),
+            None,
             [(0, (0, 0)), (1, (1, 1))],
             ["a", "v", "c"],
             [{"a": 1, "c": 3, "v": "a"}, {"a": 2, "c": 4, "v": "b"}],
         ),
         (
             ["a", "v", "c"],
+            None,
             [
                 ((0, 0), 0),
                 ((0, 0), 1),
@@ -155,6 +161,7 @@ def test_splits_groups_comb(
         ),
         (
             [["a", "v"], "c"],
+            None,
             [
                 ((0, 0), 0),
                 ((0, 0), 1),
@@ -179,6 +186,7 @@ def test_splits_groups_comb(
         ),
         (
             ["a", ["v", "c"]],
+            None,
             [
                 (0, (0, 0)),
                 (0, (0, 1)),
@@ -203,6 +211,7 @@ def test_splits_groups_comb(
         ),
         (
             [("a", "v"), "c"],
+            None,
             [((0, 0), 0), ((0, 0), 1), ((1, 1), 0), ((1, 1), 1)],
             ["a", "v", "c"],
             [
@@ -214,6 +223,7 @@ def test_splits_groups_comb(
         ),
         (
             ["a", ("v", "c")],
+            None,
             [(0, (0, 0)), (0, (1, 1)), (1, (0, 0)), (1, (1, 1))],
             ["a", "v", "c"],
             [
@@ -226,12 +236,14 @@ def test_splits_groups_comb(
         # TODO: check if it's ok
         (
             (("a", "v"), ("c", "z")),
+            None,
             [((0, 0), (0, 0)), ((1, 1), (1, 1))],
             ["a", "v", "c", "z"],
             [{"a": 1, "v": "a", "c": 3, "z": 7}, {"a": 2, "v": "b", "c": 4, "z": 8}],
         ),
         (
             (["a", "v"], ["c", "z"]),
+            None,
             [((0, 0), (0, 0)), ((0, 1), (0, 1)), ((1, 0), (1, 0)), ((1, 1), (1, 1))],
             ["a", "v", "c", "z"],
             [
@@ -243,6 +255,7 @@ def test_splits_groups_comb(
         ),
         (
             [("a", "v"), ("c", "z")],
+            None,
             [((0, 0), (0, 0)), ((0, 0), (1, 1)), ((1, 1), (0, 0)), ((1, 1), (1, 1))],
             ["a", "v", "c", "z"],
             [
@@ -254,6 +267,7 @@ def test_splits_groups_comb(
         ),
         (
             (["a", "v"], "x"),
+            {"x": 2},  # input x is treated as 2d container, so x will be flatten
             [((0, 0), 0), ((0, 1), 1), ((1, 0), 2), ((1, 1), 3)],
             ["a", "v", "x"],
             [
@@ -265,7 +279,7 @@ def test_splits_groups_comb(
         ),
     ],
 )
-def test_splits_1b(splitter, values, keys, splits):
+def test_splits_1b(splitter, cont_dim, values, keys, splits):
     inputs = {
         "a": [1, 2],
         "v": ["a", "b"],
@@ -274,40 +288,51 @@ def test_splits_1b(splitter, values, keys, splits):
         "x": [[10, 100], [20, 200]],
     }
     splitter_rpn = hlpst.splitter2rpn(splitter)
-    values_out, keys_out, _, _ = hlpst.splits(splitter_rpn, inputs)
+    values_out, keys_out, _ = hlpst.splits(splitter_rpn, inputs, cont_dim=cont_dim)
     value_list = list(values_out)
     assert keys == keys_out
     assert values == value_list
-    splits_out = list(hlpst.map_splits(hlpst.iter_splits(value_list, keys_out), inputs))
+    splits_out = list(
+        hlpst.map_splits(
+            hlpst.iter_splits(value_list, keys_out), inputs, cont_dim=cont_dim
+        )
+    )
     assert splits_out == splits
 
 
 @pytest.mark.parametrize(
-    "splitter, inputs, mismatch",
+    "splitter, cont_dim, inputs, mismatch",
     [
-        ((["a", "v"], "c"), {"a": [1, 2], "v": ["a", "b"], "c": [3, 4]}, True),
+        ((["a", "v"], "c"), None, {"a": [1, 2], "v": ["a", "b"], "c": [3, 4]}, True),
         (
             (["a", "v"], "c"),
+            {"c": 2},  # c is treated as 2d container
             {"a": [1, 2], "v": ["a", "b"], "c": [[3, 4], [5, 6]]},
             False,
         ),
-        ((["a", "v"], "c"), {"a": [1, 2], "v": ["a", "b"], "c": [[3, 4], [5]]}, True),
+        (
+            (["a", "v"], "c"),
+            None,
+            {"a": [1, 2], "v": ["a", "b"], "c": [[3, 4], [5]]},
+            True,
+        ),
     ],
 )
-def test_splits_1c(splitter, inputs, mismatch):
+def test_splits_1c(splitter, cont_dim, inputs, mismatch):
     splitter_rpn = hlpst.splitter2rpn(splitter)
     if mismatch:
         with pytest.raises(ValueError):
-            hlpst.splits(splitter_rpn, inputs)
+            hlpst.splits(splitter_rpn, inputs, cont_dim=cont_dim)
     else:
-        hlpst.splits(splitter_rpn, inputs)
+        hlpst.splits(splitter_rpn, inputs, cont_dim=cont_dim)
 
 
 @pytest.mark.parametrize(
-    "splitter, values, keys, shapes, splits",
+    "splitter, cont_dim, values, keys, shapes, splits",
     [
         (
             (["a", "v"], "c"),
+            {"c": 2},
             [((0, 0), 0), ((0, 1), 1), ((1, 0), 2), ((1, 1), 3)],
             ["a", "v", "c"],
             {"a": (2,), "v": (2,), "c": (2, 2)},
@@ -320,6 +345,7 @@ def test_splits_1c(splitter, inputs, mismatch):
         ),
         (
             ("c", ["a", "v"]),
+            {"c": 2},
             [(0, (0, 0)), (1, (0, 1)), (2, (1, 0)), (3, (1, 1))],
             ["c", "a", "v"],
             {"a": (2,), "v": (2,), "c": (2, 2)},
@@ -332,15 +358,18 @@ def test_splits_1c(splitter, inputs, mismatch):
         ),
     ],
 )
-def test_splits_1d(splitter, values, keys, shapes, splits):
+def test_splits_1d(splitter, cont_dim, values, keys, shapes, splits):
     inputs = {"a": [1, 2], "v": ["a", "b"], "c": [[3, 4], [5, 6]]}
     splitter_rpn = hlpst.splitter2rpn(splitter)
-    values_out, keys_out, shapes_out, _ = hlpst.splits(splitter_rpn, inputs)
+    values_out, keys_out, _ = hlpst.splits(splitter_rpn, inputs, cont_dim=cont_dim)
     value_list = list(values_out)
     assert keys == keys_out
     assert values == value_list
-    assert shapes == shapes_out
-    splits_out = list(hlpst.map_splits(hlpst.iter_splits(value_list, keys_out), inputs))
+    splits_out = list(
+        hlpst.map_splits(
+            hlpst.iter_splits(value_list, keys_out), inputs, cont_dim=cont_dim
+        )
+    )
     assert splits_out == splits
 
 
@@ -371,7 +400,7 @@ def test_splits_1e(splitter, values, keys, splits):
     # c - is like an inner splitter
     inputs = {"a": [1, 2], "v": ["a", "b"], "c": [[3, 4], 5]}
     splitter_rpn = hlpst.splitter2rpn(splitter)
-    values_out, keys_out, _, _ = hlpst.splits(splitter_rpn, inputs)
+    values_out, keys_out, _ = hlpst.splits(splitter_rpn, inputs)
     value_list = list(values_out)
     assert keys == keys_out
     assert values == value_list
@@ -406,18 +435,24 @@ def test_splits_2(splitter_rpn, inner_inputs, values, keys, splits):
         "NA.a": ["a1", "a2"],
         "NA.b": ["b1", "b2"],
         "NB.b": [["b11", "b12"], ["b21", "b22"]],
-        "c": ["c1", "c2"],
-        "NB.d": [
-            [["d111", "d112"], ["d121", "d122"]],
-            [["d211", "d212"], ["d221", "d222"]],
-        ],
+        # needed?
+        # "c": ["c1", "c2"],
+        # "NB.d": [
+        #     [["d111", "d112"], ["d121", "d122"]],
+        #     [["d211", "d212"], ["d221", "d222"]],
+        # ],
     }
-    values_out, keys_out, _, _ = hlpst.splits(
-        splitter_rpn, inputs, inner_inputs=inner_inputs
+    cont_dim = {"NB.b": 2}  # will be treated as 2d container
+    values_out, keys_out, _ = hlpst.splits(
+        splitter_rpn, inputs, inner_inputs=inner_inputs, cont_dim=cont_dim
     )
     value_list = list(values_out)
     assert keys == keys_out
-    splits_out = list(hlpst.map_splits(hlpst.iter_splits(value_list, keys_out), inputs))
+    splits_out = list(
+        hlpst.map_splits(
+            hlpst.iter_splits(value_list, keys_out), inputs, cont_dim=cont_dim
+        )
+    )
     assert splits_out == splits
 
 
diff --git a/pydra/engine/tests/test_node_task.py b/pydra/engine/tests/test_node_task.py
index 26aa15d4..ff49ea79 100644
--- a/pydra/engine/tests/test_node_task.py
+++ b/pydra/engine/tests/test_node_task.py
@@ -832,6 +832,27 @@ def test_task_state_2(plugin):
 
 @pytest.mark.parametrize("plugin", Plugins)
 def test_task_state_3(plugin):
+    """ task with a list as an input, and a simple splitter """
+    nn = moment(name="NA", n=3, lst=[[2, 3, 4], [1, 2, 3]]).split(splitter="lst")
+    assert np.allclose(nn.inputs.n, 3)
+    assert np.allclose(nn.inputs.lst, [[2, 3, 4], [1, 2, 3]])
+    assert nn.state.splitter == "NA.lst"
+
+    with Submitter(plugin=plugin) as sub:
+        sub(nn)
+
+    # checking the results
+    results = nn.result()
+    for i, expected in enumerate([33, 12]):
+        assert results[i].output.out == expected
+    # checking the output_dir
+    assert nn.output_dir
+    for odir in nn.output_dir:
+        assert odir.exists()
+
+
+@pytest.mark.parametrize("plugin", Plugins)
+def test_task_state_3a(plugin):
     """ task with a tuple as an input, and a simple splitter """
     nn = moment(name="NA", n=3, lst=[(2, 3, 4), (1, 2, 3)]).split(splitter="lst")
     assert np.allclose(nn.inputs.n, 3)
@@ -853,8 +874,8 @@ def test_task_state_3(plugin):
 
 @pytest.mark.parametrize("plugin", Plugins)
 def test_task_state_4(plugin):
-    """ task with a tuple as an input, and the variable is part of the scalar splitter"""
-    nn = moment(name="NA", n=[1, 3], lst=[(2, 3, 4), (1, 2, 3)]).split(
+    """ task with a list as an input, and the variable is part of the scalar splitter"""
+    nn = moment(name="NA", n=[1, 3], lst=[[2, 3, 4], [1, 2, 3]]).split(
         splitter=("n", "lst")
     )
     assert np.allclose(nn.inputs.n, [1, 3])
@@ -876,10 +897,10 @@ def test_task_state_4(plugin):
 
 @pytest.mark.parametrize("plugin", Plugins)
 def test_task_state_4_exception(plugin):
-    """ task with a tuple as an input, and the variable is part of the scalar splitter
+    """ task with a list as an input, and the variable is part of the scalar splitter
         the shapes are not matching, so exception should be raised
     """
-    nn = moment(name="NA", n=[1, 3, 3], lst=[(2, 3, 4), (1, 2, 3)]).split(
+    nn = moment(name="NA", n=[1, 3, 3], lst=[[2, 3, 4], [1, 2, 3]]).split(
         splitter=("n", "lst")
     )
     assert np.allclose(nn.inputs.n, [1, 3, 3])
@@ -894,6 +915,29 @@ def test_task_state_4_exception(plugin):
 
 @pytest.mark.parametrize("plugin", Plugins)
 def test_task_state_5(plugin):
+    """ ask with a list as an input, and the variable is part of the outer splitter """
+    nn = moment(name="NA", n=[1, 3], lst=[[2, 3, 4], [1, 2, 3]]).split(
+        splitter=["n", "lst"]
+    )
+    assert np.allclose(nn.inputs.n, [1, 3])
+    assert np.allclose(nn.inputs.lst, [[2, 3, 4], [1, 2, 3]])
+    assert nn.state.splitter == ["NA.n", "NA.lst"]
+
+    with Submitter(plugin=plugin) as sub:
+        sub(nn)
+
+    # checking the results
+    results = nn.result()
+    for i, expected in enumerate([3, 2, 33, 12]):
+        assert results[i].output.out == expected
+    # checking the output_dir
+    assert nn.output_dir
+    for odir in nn.output_dir:
+        assert odir.exists()
+
+
+@pytest.mark.parametrize("plugin", Plugins)
+def test_task_state_5a(plugin):
     """ ask with a tuple as an input, and the variable is part of the outer splitter """
     nn = moment(name="NA", n=[1, 3], lst=[(2, 3, 4), (1, 2, 3)]).split(
         splitter=["n", "lst"]
diff --git a/pydra/engine/tests/test_state.py b/pydra/engine/tests/test_state.py
index a42cae49..49100394 100644
--- a/pydra/engine/tests/test_state.py
+++ b/pydra/engine/tests/test_state.py
@@ -202,7 +202,7 @@ def test_state_connect_2():
 
 def test_state_connect_2a():
     """ two 'connected' states: testing groups, prepare_states and prepare_inputs
-         the second state has explicit splitter that contains 
+         the second state has explicit splitter that contains
          splitter from the first node and a new field;
          adding an additional scalar field that is not part of the splitter
     """
@@ -564,7 +564,10 @@ def test_state_connect_innerspl_1():
     assert st2.group_for_inputs_final == {"NA.a": 0, "NB.b": 1}
     assert st2.groups_stack_final == [[0], [1]]
 
-    st2.prepare_states(inputs={"NA.a": [3, 5], "NB.b": [[1, 10, 100], [2, 20, 200]]})
+    st2.prepare_states(
+        inputs={"NA.a": [3, 5], "NB.b": [[1, 10, 100], [2, 20, 200]]},
+        cont_dim={"NB.b": 2},  # will be treated as 2d container
+    )
 
     assert st2.states_ind == [
         {"NA.a": 0, "NB.b": 0},
@@ -608,7 +611,10 @@ def test_state_connect_innerspl_1a():
     assert st2.group_for_inputs_final == {"NA.a": 0, "NB.b": 1}
     assert st2.groups_stack_final == [[0], [1]]
 
-    st2.prepare_states(inputs={"NA.a": [3, 5], "NB.b": [[1, 10, 100], [2, 20, 200]]})
+    st2.prepare_states(
+        inputs={"NA.a": [3, 5], "NB.b": [[1, 10, 100], [2, 20, 200]]},
+        cont_dim={"NB.b": 2},  # will be treated as 2d container
+    )
 
     assert st2.states_ind == [
         {"NA.a": 0, "NB.b": 0},
@@ -660,7 +666,8 @@ def test_state_connect_innerspl_2():
     assert st2.groups_stack_final == [[0], [1, 2]]
 
     st2.prepare_states(
-        inputs={"NA.a": [3, 5], "NB.b": [[1, 10, 100], [2, 20, 200]], "NB.c": [13, 17]}
+        inputs={"NA.a": [3, 5], "NB.b": [[1, 10, 100], [2, 20, 200]], "NB.c": [13, 17]},
+        cont_dim={"NB.b": 2},  # will be treated as 2d container
     )
     assert st2.states_ind == [
         {"NB.c": 0, "NA.a": 0, "NB.b": 0},
@@ -724,7 +731,8 @@ def test_state_connect_innerspl_2a():
     assert st2.groups_stack_final == [[0], [1, 2]]
 
     st2.prepare_states(
-        inputs={"NA.a": [3, 5], "NB.b": [[1, 10, 100], [2, 20, 200]], "NB.c": [13, 17]}
+        inputs={"NA.a": [3, 5], "NB.b": [[1, 10, 100], [2, 20, 200]], "NB.c": [13, 17]},
+        cont_dim={"NB.b": 2},  # will be treated as 2d container
     )
 
     assert st2.states_ind == [
@@ -794,7 +802,8 @@ def test_state_connect_innerspl_3():
             "NB.b": [[1, 10, 100], [2, 20, 200]],
             "NB.c": [13, 17],
             "NC.d": [33, 77],
-        }
+        },
+        cont_dim={"NB.b": 2},  # will be treated as 2d container
     )
 
     assert st2.states_ind == [
@@ -933,7 +942,8 @@ def test_state_connect_innerspl_4():
             "NC.e": [30, 50],
             "NC.f": [[23, 27], [33, 37]],
             "NC.d": [1, 2],
-        }
+        },
+        cont_dim={"NC.f": 2},  # will be treated as 2d container
     )
     assert st3.states_ind == [
         {"NA.a": 0, "NB.b": 0, "NB.c": 0, "NC.d": 0},
@@ -1198,7 +1208,8 @@ def test_state_connect_innerspl_combine_1():
     # assert st2.groups_stack_final == [[0, 1]]
 
     st2.prepare_states(
-        inputs={"NA.a": [3, 5], "NB.b": [[1, 10, 100], [2, 20, 200]], "NB.c": [13, 17]}
+        inputs={"NA.a": [3, 5], "NB.b": [[1, 10, 100], [2, 20, 200]], "NB.c": [13, 17]},
+        cont_dim={"NB.b": 2},  # will be treated as 2d container
     )
     # NOW TODO: checking st2.states_ind_final!!!
     assert st2.states_ind == [
@@ -1271,7 +1282,8 @@ def test_state_connect_innerspl_combine_2():
     assert st2.groups_stack_final == [[0], [1]]
 
     st2.prepare_states(
-        inputs={"NA.a": [3, 5], "NB.b": [[1, 10, 100], [2, 20, 200]], "NB.c": [13, 17]}
+        inputs={"NA.a": [3, 5], "NB.b": [[1, 10, 100], [2, 20, 200]], "NB.c": [13, 17]},
+        cont_dim={"NB.b": 2},  # will be treated as 2d container
     )
     assert st2.states_ind == [
         {"NB.c": 0, "NA.a": 0, "NB.b": 0},