keras/src/ops/numpy.py/TakeAlongAxis/compute_output_spec
class TakeAlongAxis:    def compute_output_spec(self, x, indices):
        output_shape = operation_utils.compute_take_along_axis_output_shape(
            x.shape, indices.shape, self.axis
        )
        return KerasTensor(output_shape, dtype=x.dtype)