keras/src/ops/numpy.py/Average/compute_output_spec
class Average:    def compute_output_spec(self, x, weights=None):
        dtypes_to_resolve = [getattr(x, "dtype", type(x)), float]
        if weights is not None:
            shape_match = shape_equal(x.shape, weights.shape, allow_none=True)
            if self.axis is not None:
                shape_match_on_axis = shape_equal(
                    [x.shape[self.axis]], weights.shape, allow_none=True
                )
            dtypes_to_resolve.append(getattr(weights, "dtype", type(weights)))
        dtype = dtypes.result_type(*dtypes_to_resolve)
        if self.axis is None:
            if weights is None or shape_match:
                return KerasTensor([], dtype=dtype)
            else:
                raise ValueError(
                    "`weights` must have the same shape as `x` when "
                    f"`axis=None`, but received `weights.shape={weights.shape}`"
                    f" and `x.shape={x.shape}`."
                )

        if weights is None or shape_match_on_axis or shape_match:
            return KerasTensor(
                reduce_shape(x.shape, axis=[self.axis]), dtype=dtype
            )
        else:
            # `weights` can either be a 1D array of length `x.shape[axis]` or
            # of the same shape as `x`.
            raise ValueError(
                "`weights` must have the same size as `x` at "
                f"`axis={self.axis}` but received "
                f"`weights.shape={weights.shape}` while x.shape at "
                f"`{self.axis}` is `{x.shape[self.axis]}`."
            )