How to determine next token indices and slice output tensor in concrete function utilizing model prediction via model(prompt)
Hi **@andrew-lyons** ,

Could you consider techniques like:Optimizing tensor slicing operations.
Using XLA (Accelerated Linear Algebra) compilation for faster computations. I hope it will be useful to you.

Thank you! This issue is stale because it has been open for 7 days with no activity. It will be closed if no further activity occurs. Thank you. This issue was closed because it has been inactive for 7 days since being marked as stale. Please reopen if you'd like to work on this further. Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/62677">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/62677">No</a>

### 1. System information

- OS Platform and Distribution (e.g., Linux Ubuntu 16.04):
  - Mac, so running this via Google Colab
- TensorFlow installation (pip package or built from source):
  - built via `pip install https://github.com/keras-team/keras-nlp.git@google-io-2023`
- TensorFlow library (version, if pip package or github SHA, if built from source):
  - `tensorflow-text==2.12`
### 2. Code

Provide code to help us reproduce your issues using one of the following options:

#### Option B: Paste your code here or provide a link to a custom end-to-end colab

```
import numpy as np
import keras_nlp
import tensorflow as tf
import tensorflow_datasets as tfds
import tensorflow_text as tf_text
from tensorflow import keras
from tensorflow.lite.python import interpreter
import time

gpt2_tokenizer = keras_nlp.models.GPT2Tokenizer.from_preset("gpt2_base_en")
gpt2_preprocessor = keras_nlp.models.GPT2CausalLMPreprocessor.from_preset(
    "gpt2_base_en",
    sequence_length=256,
    add_end_token=True,
)
gpt2_lm = keras_nlp.models.GPT2CausalLM.from_preset("gpt2_base_en", preprocessor=gpt2_preprocessor)

@tf.function
def predict(prompt):
  preprocessed = gpt2_preprocessor(prompt)[0]
  input_map = {
      'token_ids': tf.reshape(preprocessed['token_ids'], [1, 255]),
      'padding_mask': tf.reshape(preprocessed['padding_mask'], [1, 255])
  }

  result = gpt2_lm(input_map)[0] # shape of (1, 255, 50256)

  # ideally slice output here?

  topk = tf.math.top_k(tf.nn.softmax(result), k=25)
  return topk.indices

concrete_func = predict.get_concrete_function(tf.TensorSpec([], tf.string))
```

I haven't included my export/conversion to TFLite but everything there is working fine. In the concrete function I'm topk sampling (25) which means my output of this function is 255*25 tokens long.

My question is to how I could understand where the next token lies in this 255 output sequence and only return that, making my output an int32[] of size 25?

If I was to do this outside of the `tf.function`, I would manually tokenize the input prompt, get the shape of the resulting tensor, and slice the model's output using that shape. However I understand that I cannot due this inside the `tf.function`. Is it possible to perform the same kind of logic using tf builtins?