code: def _split_state(
      state: State,
      *filters: filterlib.Filter,
  ) -> tp.Tuple[State, ...]:
      for i, filter_ in enumerate(filters):
          if filter_ is ... and i != len(filters) - 1:
              raise ValueError(
                  "Ellipsis `...` can only be used as the last filter, "
                  f"got it at index {i}."
              )
      predicates = tuple(map(filterlib.to_predicate, filters))
  
      flat_state = state.flat_state()
  
      # we have n + 1 states, where n is the number of predicates
      # the last state is for values that don't match any predicate
      flat_states: tp.Tuple[FlatState, ...] = tuple(
          {} for _ in range(len(predicates) + 1)
      )
  
      for path, value in flat_state.items():
          for i, predicate in enumerate(predicates):
              if predicate(path, value):
                  flat_states[i][path] = value
                  break
          else:
              # if we didn't break, set leaf to last state
              flat_states[-1][path] = value
  
      return tuple(
          State.from_flat_path(flat_state) for flat_state in flat_states
      )
  
imports: from __future__ import annotations
  import typing
  import jax
  import jax.tree_util
  from flax import traverse_util
  from flax.experimental.nnx.nnx import filterlib, reprlib
  from flax.experimental.nnx.nnx.variables import Variable
label_desc: Use image node_id 4 for calling a global function with example usage: _split_state(state) and returns: tuple