Predict crashes if last batch is smaller then batch_size
Hi @AdnanAvdagic ,

Could you please submit minimal reproducible code snippet for this. I tried with a dummy model with an input size that is not a multiple of batch_size but couldn't get any issue.Attached [gist](https://colab.sandbox.google.com/gist/SuryanarayanaY/231d7d802afc70c568611129ae40b7f2/63811_sample_model-tf2-16v.ipynb) for reference. ```python
backbone = keras_cv.models.YOLOV8Backbone.from_preset(
                'yolo_v8_l_backbone',
                input_shape=[512,512, 3],
            )

model = keras_cv.models.YOLOV8Detector(
                    num_classes=1,
                    bounding_box_format='xyxy',
                    backbone=backbone,
                    fpn_depth=3,
                )

model.prediction_decoder = keras_cv.layers.NonMaxSuppression(
            bounding_box_format='xyxy',
            from_logits=True,
            iou_threshold=0.35,
            confidence_threshold=0.53,
            max_detections=50,
        )

def prepare_image(self: Self, input_image_path: Path) -> np.ndarray:
        try:
            input_image = cv2.imread(str(input_image_path))
            # input_image = cv2.bilateralFilter(input_image, 9, 75, 75)
            if self.color_mode == "rgb":
                input_image = cv2.cvtColor(input_image, cv2.COLOR_BGR2RGB)

            else:
                input_image = cv2.cvtColor(input_image, cv2.COLOR_BGR2GRAY)

            return input_image
tiles_paths = ['01.jpg', '02.jpg', ...]

image_patches = np.array(
  [
    prepare_image(input_image_path=tile_path)
                        for tile_path in tiles_paths
  ],
)
y_pred = model.predict(
   image_patches,
   batch_size=8,
)
```

Here we get a crash if the number of image_patches is not divisible with the batch_size it will crash at the last batch which is less than the batch_size Hi @AdnanAvdagic ,

Since this issue seems specific to `keras_cv` model. Could you please file an issue at [keras_cv](https://github.com/keras-team/keras-cv/issues) repo,since all `keras_cv` issues are actively tracked there? 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/63811">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/63811">No</a>

### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

No

### Source

source

### TensorFlow version

tf 2.16.1

### Custom code

Yes

### OS platform and distribution

Linux Ubuntu 23.04

### Mobile device

_No response_

### Python version

3.11

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

CUDA 12.4 cuDNN 8.9

### GPU model and memory

RTX 4090

### Current behavior?

When sending in a numpy array which is not exactly divisible into the batch size, ex length of numpy array is 100 and batch size is 8 the inferencing crashes.

In the function take_along_axis we found that the batch_size axis was None instead of 8 for the last batch

### Standalone code to reproduce the issue

```python
image_patches = np.array([self.prepare_image(tile_path) for tile_path in tiles_paths])
y_pred = model.predict(
    image_patches,
    batch_size=8
)
```


### Relevant log output

```shell
228/229 ━━━━━━━━━━━━━━━━━━━━ 0s 74ms/step2024-03-15 16:02:37.436169: W tensorflow/core/framework/op_kernel.cc:1839] OP_REQUIRES failed at bcast_ops.cc:52 : INVALID_ARGUMENT: Input 0 to node non_max_suppression_1_1/BroadcastArgs with op BroadcastArgs must be a compile-time constant.

XLA compilation requires that operator arguments that represent shapes or dimensions be evaluated to concrete values at compile time. This error means that a shape or dimension argument could not be evaluated at compile time, usually because the value of the argument depends on a parameter to the computation, on a variable, or on a stateful operation such as a random number generator.

Stack trace for op definition: 
File "main_inferencer.py", line 231, in <module>
File "utils/exception_utils/inspectors.py", line 317, in wrapper
File "main_inferencer.py", line 130, in main
File "src/e_inference/src/geotif_inferencer.py", line 59, in main
File "src/e_inference/src/geotif_inferencer.py", line 168, in detect_image
File ".venv/lib/python3.11/site-packages/keras/src/utils/traceback_utils.py", line 118, in error_handler
File ".venv/lib/python3.11/site-packages/keras/src/backend/tensorflow/trainer.py", line 513, in predict
File ".venv/lib/python3.11/site-packages/keras/src/backend/tensorflow/trainer.py", line 212, in one_step_on_data_distributed
File ".venv/lib/python3.11/site-packages/keras/src/backend/tensorflow/trainer.py", line 201, in one_step_on_data
File ".venv/lib/python3.11/site-packages/keras_cv/src/models/object_detection/yolo_v8/yolo_v8_detector.py", line 616, in predict_step
File ".venv/lib/python3.11/site-packages/keras_cv/src/models/object_detection/yolo_v8/yolo_v8_detector.py", line 609, in decode_predictions
File ".venv/lib/python3.11/site-packages/keras/src/utils/traceback_utils.py", line 118, in error_handler
File ".venv/lib/python3.11/site-packages/keras/src/layers/layer.py", line 816, in __call__
File ".venv/lib/python3.11/site-packages/keras/src/utils/traceback_utils.py", line 118, in error_handler
File ".venv/lib/python3.11/site-packages/keras/src/ops/operation.py", line 42, in __call__
File ".venv/lib/python3.11/site-packages/keras/src/utils/traceback_utils.py", line 157, in error_handler
File ".venv/lib/python3.11/site-packages/keras_cv/src/layers/object_detection/non_max_suppression.py", line 143, in call
File ".venv/lib/python3.11/site-packages/keras/src/ops/numpy.py", line 4991, in take_along_axis
File ".venv/lib/python3.11/site-packages/keras/src/backend/tensorflow/numpy.py", line 1505, in take_along_axis
```
