keras/src/ops/numpy.py/Take/compute_output_spec
class Take:    def compute_output_spec(self, x, indices):
        x_shape = list(x.shape)
        if isinstance(indices, KerasTensor):
            indices_shape = list(indices.shape)
        else:
            indices_shape = list(getattr(np.array(indices), "shape", []))
        if self.axis is None:
            return KerasTensor(indices_shape, dtype=x.dtype)

        # make sure axis is non-negative
        axis = len(x_shape) + self.axis if self.axis < 0 else self.axis
        output_shape = x_shape[:axis] + indices_shape + x_shape[axis + 1 :]
        return KerasTensor(output_shape, dtype=x.dtype)