Function init_decode_state

src/MaxText/maxtext_utils.py:743–746

Init train state with null opt state for decode.

(apply_fn, params)

def init_decode_state(apply_fn, params) -> train_state.TrainState:
  """Init train state with null opt state for decode."""
  state = train_state.TrainState(step=0, apply_fn=apply_fn, params=params, tx=None, opt_state={})  # type: ignore
  return state


def init_training_state(apply_fn, params, tx):