def METHOD_NAME(
    train_state: TrainState,
    checkpoint_path: str | None = None,
    step: int | None = None,
) -> TrainStateProvenance:
  assert not isinstance(
      train_state.step, jax.sharding.PartitionSpec
  ), 'Tensor provenance is only for tensors'
  provenance = TensorProvenance()
  if checkpoint_path:
    provenance = TensorProvenance(
        checkpoint_path=checkpoint_path, checkpoint_step=step
    )
  return TrainStateProvenance(
      step=provenance,
      mdl_vars=jax.tree_map(lambda x: provenance, train_state.mdl_vars),
      opt_states=jax.tree_map(lambda x: provenance, train_state.opt_states),
      extra_state=jax.tree_map(lambda x: provenance, train_state.extra_state),
  )