def update_state(self, y_true, y_pred, sample_weight=None):
        """Accumulates root mean squared error statistics.

        Args:
            y_true: The ground truth values.
            y_pred: The predicted values.
            sample_weight: Optional weighting of each example. Can
                be a `Tensor` whose rank is either 0, or the same rank as
                `y_true`, and must be broadcastable to `y_true`.
                Defaults to `1`.

        Returns:
            Update op.
        """
        y_true = ops.convert_to_tensor(y_true, dtype=self._dtype)
        y_pred = ops.convert_to_tensor(y_pred, dtype=self._dtype)
        y_true, y_pred = squeeze_or_expand_to_same_rank(y_true, y_pred)
        if not self._built:
            self._build(y_true.shape, y_pred.shape)

        if sample_weight is None:
            sample_weight = 1

        sample_weight = ops.convert_to_tensor(sample_weight, dtype=self.dtype)

        if len(sample_weight.shape) == 1:
            # Make sure there's a features dimension
            sample_weight = ops.expand_dims(sample_weight, axis=1)

        sample_weight = ops.broadcast_to(sample_weight, ops.shape(y_true))

        weighted_y_true = y_true * ops.cast(sample_weight, y_true.dtype)
        self.sum.assign(self.sum + ops.sum(weighted_y_true, axis=0))
        self.squared_sum.assign(
            self.squared_sum + ops.sum(y_true * weighted_y_true, axis=0)
        )
        self.total_mse.assign(
            self.total_mse
            + ops.sum(
                (y_true - y_pred) ** 2 * ops.cast(sample_weight, y_true.dtype),
                axis=0,
            )
        )
        self.count.assign(self.count + ops.sum(sample_weight, axis=0))
        self.num_samples.assign(self.num_samples + ops.size(y_true))
Accumulates root mean squared error statistics.

        Args:
            y_true: The ground truth values.
            y_pred: The predicted values.
            sample_weight: Optional weighting of each example. Can
                be a `Tensor` whose rank is either 0, or the same rank as
                `y_true`, and must be broadcastable to `y_true`.
                Defaults to `1`.

        Returns:
            Update op.
        
keras/src/metrics/regression_metrics.py