def call_fn_with_state_keys(jit_fn, state, other_inputs, keys):
  """Executes `jit_fn`, filtering out all keys except some subset."""
  state = state.copy()
  extra_state = {}
  for k in list(state.keys()):
    if k not in keys:
      extra_state[k] = state.pop(k)
  return jit_fn(state, *other_inputs), extra_state