Упор AutoFocus срабатывает только при установленном компоненте. Когда вы возвращаетесь к A, если A все еще смонтирован (просто скрыт), то автофокус не будет работать снова.
Вы должны использовать ссылку (добавить новое состояние, например, ref
здесь) и добавьте обработчик при переходе назад от B к A, чтобы запустить this.state.ref.focus()
<TextInput
ref={ref => ref && this.setState({ref})}
style={styles.textInput}
textStyle={styles.inputTextStyle}
autoFocus={true}
placeholder='Enter Code'
keyboard={'numeric'}
placeholderTextColor={AppStyles.color.grey}
value={code}
onChangeText={this.updateCode}
underline={false}
/>