def _choose_states(F, states, state_info, indices):
assert not state_info or isinstance(state_info, (type(states), dict)),'<STR_LIT>'<EOL>if isinstance(states, list):<EOL><INDENT>if not state_info:<EOL><INDENT>state_info = [None] * len(states)<EOL><DEDENT>return [_choose_states(F, d, s, indices) for d, s in zip(states, state_info)]<EOL><DEDENT>elif isinstance(states, tuple):<EOL><INDENT>if not state_info:<EOL><INDENT>state_info = [None] * len(states)<EOL>state_info = tuple(state_info)<EOL><DEDENT>return tuple(_choose_states(F, d, s, indices) for d, s in zip(states, state_info))<EOL><DEDENT>elif isinstance(states, dict):<EOL><INDENT>if not state_info:<EOL><INDENT>state_info = {k: None for k in states.keys()}<EOL><DEDENT>return {k: _choose_states(F, v, state_info[k], indices)<EOL>for k, v in states.items()}<EOL><DEDENT>elif isinstance(states, (mx.nd.NDArray, mx.sym.Symbol)):<EOL><INDENT>if not state_info:<EOL><INDENT>batch_axis = <NUM_LIT:0><EOL><DEDENT>else:<EOL><INDENT>batch_axis = state_info['<STR_LIT>'].find('<STR_LIT:N>')<EOL><DEDENT>states = F.take(states, indices, axis=batch_axis)<EOL>return states<EOL><DEDENT>else:<EOL><INDENT>raise NotImplementedError<EOL><DEDENT>
Parameters
----------
F : ndarray or symbol
states : Object contains NDArrays/Symbols
    Each NDArray/Symbol should have shape (N, ...) when state_info is None,
    or same as the layout in state_info when it's not None.
state_info : Nested structure of dictionary, default None.
    Descriptors for states, usually from decoder's ``state_info()``.
    When None, this method assumes that the batch axis is the first dimension.
indices : NDArray or Symbol
    Indices of the states to take. Shape (N,).
Returns
-------
new_states : Object contains NDArrays/Symbols
    Each NDArray/Symbol should have shape (N, ...).