keras-nightly 3.14.0.dev2026010104__py3-none-any.whl → 3.14.0.dev2026012204__py3-none-any.whl

keras/src/layers/core/dense.py

@@ -128,7 +128,7 @@ class Dense(Layer):
                 mode=self.quantization_mode,
                 config=self.quantization_config,
             )
-        if self.quantization_mode not in ("int8", "int4", "gptq"):
+        if self.quantization_mode not in ("int8", "int4", "gptq", "awq"):
             # If the layer is quantized to int8 or int4, `self._kernel` will be
             # added in `self._int8_build` or `_int4_build`. Therefore, we skip
             # it here.
@@ -165,15 +165,17 @@ class Dense(Layer):
 
         mode = self.quantization_mode
         is_gptq = mode == "gptq"
+        is_awq = mode == "awq"
         is_int4 = mode == "int4"
-        calibrated = bool(getattr(self, "is_gptq_calibrated", False))
+        gptq_calibrated = bool(getattr(self, "is_gptq_calibrated", False))
+        awq_calibrated = bool(getattr(self, "is_awq_calibrated", False))
         gptq_bits = (
             gptq_core.get_weight_bits_for_layer(self, None) if is_gptq else None
         )
 
         # Decide the source tensor first (packed vs already-quantized vs plain
         # kernel)
-        if is_gptq and calibrated and gptq_bits != 4:
+        if is_gptq and gptq_calibrated and gptq_bits != 4:
             # calibrated GPTQ, not 4-bit, no unpacking needed
             kernel = self.quantized_kernel
         else:
@@ -183,7 +185,15 @@ class Dense(Layer):
             # Handle int4 unpacking cases in one place
             if is_int4:
                 kernel = quantizers.unpack_int4(kernel, self._orig_input_dim)
-            elif is_gptq and calibrated and gptq_bits == 4:
+            elif is_gptq and gptq_calibrated and gptq_bits == 4:
+                kernel = quantizers.unpack_int4(
+                    self.quantized_kernel,
+                    orig_len=self.units,
+                    axis=0,
+                    dtype="uint8",
+                )
+            elif is_awq and awq_calibrated:
+                # AWQ always uses 4-bit quantization
                 kernel = quantizers.unpack_int4(
                     self.quantized_kernel,
                     orig_len=self.units,
@@ -304,8 +314,9 @@ class Dense(Layer):
         if mode not in self.variable_serialization_spec:
             raise self._quantization_mode_error(mode)
 
-        # A saved GPTQ quantized model will always be calibrated.
+        # A saved GPTQ/AWQ quantized model will always be calibrated.
         self.is_gptq_calibrated = mode == "gptq"
+        self.is_awq_calibrated = mode == "awq"
 
         idx = 0
         for name in self.variable_serialization_spec[mode]:
@@ -395,6 +406,14 @@ class Dense(Layer):
                 "kernel_zero",
                 "g_idx",
             ],
+            "awq": [
+                "bias",
+                "quantized_kernel",
+                "kernel_scale",
+                "kernel_zero",
+                "awq_scales",
+                "g_idx",
+            ],
         }
 
     def quantized_build(self, kernel_shape, mode, config=None):
@@ -406,6 +425,8 @@ class Dense(Layer):
             self._float8_build()
         elif mode == "gptq":
             self._gptq_build(kernel_shape, config)
+        elif mode == "awq":
+            self._awq_build(kernel_shape, config)
         else:
             raise self._quantization_mode_error(mode)
         self._is_quantized = True
@@ -515,6 +536,97 @@ class Dense(Layer):
             y = self.activation(y)
         return y
 
+    def _awq_build(self, kernel_shape, config):
+        """Build variables for AWQ quantization.
+
+        AWQ uses 4-bit quantization with per-channel AWQ scales that protect
+        salient weights based on activation magnitudes.
+        """
+        from keras.src.quantizers import awq_core
+
+        # Ensures the forward pass uses the original high-precision kernel
+        # until calibration has been performed.
+        self.is_awq_calibrated = False
+        self.kernel_shape = kernel_shape
+
+        # For 4-bit weights, we pack two values per byte.
+        units = (kernel_shape[1] + 1) // 2
+
+        self.quantized_kernel = self.add_weight(
+            name="kernel",
+            shape=(units, kernel_shape[0]),
+            initializer="zeros",
+            dtype="uint8",
+            trainable=False,
+        )
+
+        group_size = awq_core.get_group_size_for_layer(self, config)
+        num_groups = (
+            1 if group_size == -1 else math.ceil(kernel_shape[0] / group_size)
+        )
+        self.kernel_scale = self.add_weight(
+            name="kernel_scale",
+            shape=(self.units, num_groups),
+            initializer="ones",
+            trainable=False,
+        )
+        self.kernel_zero = self.add_weight(
+            name="kernel_zero",
+            shape=(self.units, num_groups),
+            initializer="zeros",
+            dtype="uint8",
+            trainable=False,
+        )
+
+        # Per-channel AWQ scales from activation magnitudes
+        self.awq_scales = self.add_weight(
+            name="awq_scales",
+            shape=(kernel_shape[0],),
+            initializer="ones",
+            trainable=False,
+        )
+        self.g_idx = self.add_weight(
+            name="g_idx",
+            shape=(kernel_shape[0],),
+            initializer="zeros",
+            dtype="float32",
+            trainable=False,
+        )
+
+    def _awq_call(self, inputs, training=False):
+        """Forward pass for AWQ quantized layer."""
+        if not self.is_awq_calibrated:
+            W = self._kernel
+        else:
+            # Unpack 4-bit weights
+            W = quantizers.unpack_int4(
+                self.quantized_kernel,
+                orig_len=self.units,
+                axis=0,
+                dtype="uint8",
+            )
+            # Dequantize using scale/zero maps
+            W = ops.transpose(
+                dequantize_with_sz_map(
+                    W,
+                    self.kernel_scale,
+                    self.kernel_zero,
+                    self.g_idx,
+                )
+            )
+            # Apply AWQ scales by dividing to restore original magnitude
+            # (We multiplied by scales before quantization, so divide to undo)
+            # awq_scales has shape [input_dim], W has shape [input_dim, units]
+            # Expand dims for proper broadcasting.
+            W = ops.divide(W, ops.expand_dims(self.awq_scales, -1))
+
+        y = ops.matmul(inputs, W)
+        if self.bias is not None:
+            y = ops.add(y, self.bias)
+        if self.activation is not None:
+            y = self.activation(y)
+        return y
+
     def _int4_build(self, kernel_shape, config=None):
         """Build variables for int4 quantization.
 
@@ -835,6 +947,8 @@ class Dense(Layer):
             self.kernel_scale.assign(kernel_scale)
         elif mode == "gptq":
             self.quantized_build(kernel_shape, mode, self.quantization_config)
+        elif mode == "awq":
+            self.quantized_build(kernel_shape, mode, self.quantization_config)
         elif mode == "float8":
             self.quantized_build(kernel_shape, mode)
         else:
@@ -847,6 +961,8 @@ class Dense(Layer):
             policy_name = mode
             if mode == "gptq":
                 policy_name = self.quantization_config.dtype_policy_string()
+            elif mode == "awq":
+                policy_name = self.quantization_config.dtype_policy_string()
             policy = dtype_policies.get(
                 f"{policy_name}_from_{self.dtype_policy.name}"
             )
@@ -881,7 +997,7 @@ class Dense(Layer):
                 `kernel_scale`: The quantization scale for the merged kernel.
                     This is `None` if the layer is not quantized.
         """
-        if self.dtype_policy.quantization_mode in (None, "gptq"):
+        if self.dtype_policy.quantization_mode in (None, "gptq", "awq"):
             return self.kernel, None
 
         kernel_value = self._kernel

keras/src/layers/core/einsum_dense.py

@@ -180,7 +180,7 @@ class EinsumDense(Layer):
         # quantized to int8 or int4, because `quantized_build` has created the
         # appropriate kernel variable. For other modes (e.g., float8 or no
         # quantization), we still need the floating-point kernel.
-        if self.quantization_mode not in ("int8", "int4", "gptq"):
+        if self.quantization_mode not in ("int8", "int4", "gptq", "awq"):
             # If the layer is quantized to int8, `self._kernel` will be added
             # in `self._int8_build`. Therefore, we skip it here.
             self._kernel = self.add_weight(
@@ -219,15 +219,17 @@ class EinsumDense(Layer):
 
         mode = self.quantization_mode
         is_gptq = mode == "gptq"
+        is_awq = mode == "awq"
         is_int4 = mode == "int4"
-        calibrated = bool(getattr(self, "is_gptq_calibrated", False))
+        gptq_calibrated = bool(getattr(self, "is_gptq_calibrated", False))
+        awq_calibrated = bool(getattr(self, "is_awq_calibrated", False))
         gptq_bits = (
             gptq_core.get_weight_bits_for_layer(self, None) if is_gptq else None
         )
 
         # Decide the source tensor first (packed vs already-quantized vs plain
         # kernel)
-        if is_gptq and calibrated and gptq_bits != 4:
+        if is_gptq and gptq_calibrated and gptq_bits != 4:
             # calibrated GPTQ, not 4-bit, no unpacking needed
             kernel = self.quantized_kernel
         else:
@@ -241,13 +243,21 @@ class EinsumDense(Layer):
                     self._orig_length_along_pack_axis,
                     self._int4_pack_axis,
                 )
-            elif is_gptq and calibrated and gptq_bits == 4:
+            elif is_gptq and gptq_calibrated and gptq_bits == 4:
                 kernel = quantizers.unpack_int4(
                     self.quantized_kernel,
                     orig_len=self.gptq_unpacked_column_size,
                     axis=0,
                     dtype="uint8",
                 )
+            elif is_awq and awq_calibrated:
+                # AWQ always uses 4-bit quantization
+                kernel = quantizers.unpack_int4(
+                    self.quantized_kernel,
+                    orig_len=self.awq_unpacked_column_size,
+                    axis=0,
+                    dtype="uint8",
+                )
 
         # Apply LoRA if enabled
         if self.lora_enabled:
@@ -362,8 +372,9 @@ class EinsumDense(Layer):
         if mode not in self.variable_serialization_spec:
             raise self._quantization_mode_error(mode)
 
-        # A saved GPTQ quantized model will always be calibrated.
+        # A saved GPTQ/AWQ quantized model will always be calibrated.
         self.is_gptq_calibrated = mode == "gptq"
+        self.is_awq_calibrated = mode == "awq"
 
         idx = 0
         for name in self.variable_serialization_spec[mode]:
@@ -459,6 +470,14 @@ class EinsumDense(Layer):
                 "kernel_zero",
                 "g_idx",
             ],
+            "awq": [
+                "bias",
+                "quantized_kernel",
+                "kernel_scale",
+                "kernel_zero",
+                "awq_scales",
+                "g_idx",
+            ],
         }
 
     def quantized_build(self, kernel_shape, mode, config=None):
@@ -470,6 +489,8 @@ class EinsumDense(Layer):
             self._float8_build()
         elif mode == "gptq":
             self._gptq_build(kernel_shape, config)
+        elif mode == "awq":
+            self._awq_build(kernel_shape, config)
         else:
             raise self._quantization_mode_error(mode)
         self._is_quantized = True
@@ -616,6 +637,127 @@ class EinsumDense(Layer):
             y = self.activation(y)
         return y
 
+    def _awq_build(self, kernel_shape, config):
+        """Build variables for AWQ quantization.
+
+        AWQ uses 4-bit quantization with per-channel AWQ scales that protect
+        salient weights based on activation magnitudes.
+        """
+        from keras.src.quantizers import awq_core
+
+        # Ensures the forward pass uses the original high-precision kernel
+        # until calibration has been performed.
+        self.is_awq_calibrated = False
+
+        self.original_kernel_shape = kernel_shape
+        if len(kernel_shape) == 2:
+            rows = kernel_shape[0]
+            columns = kernel_shape[1]
+        elif len(kernel_shape) == 3:
+            shape = list(self.original_kernel_shape)
+            d_model_dim_index = shape.index(max(shape))
+
+            if d_model_dim_index == 0:  # QKV projection case
+                in_features, heads, head_dim = shape
+                rows, columns = (
+                    in_features,
+                    heads * head_dim,
+                )
+            elif d_model_dim_index in [1, 2]:  # Attention Output case
+                heads, head_dim, out_features = shape
+                rows, columns = (
+                    heads * head_dim,
+                    out_features,
+                )
+            else:
+                raise ValueError("Could not determine row/column split.")
+        else:
+            raise ValueError("AWQ quantization only supports 2D or 3D kernels.")
+
+        group_size = awq_core.get_group_size_for_layer(self, config)
+        num_groups = 1 if group_size == -1 else math.ceil(rows / group_size)
+
+        self.awq_unpacked_column_size = columns
+
+        # For 4-bit weights, we pack two values per byte.
+        kernel_columns = (columns + 1) // 2
+
+        self._set_quantization_info()
+
+        self.quantized_kernel = self.add_weight(
+            name="kernel",
+            shape=(kernel_columns, rows),
+            initializer="zeros",
+            dtype="uint8",
+            trainable=False,
+        )
+
+        self.kernel_scale = self.add_weight(
+            name="kernel_scale",
+            shape=(columns, num_groups),
+            initializer="ones",
+            trainable=False,
+        )
+        self.kernel_zero = self.add_weight(
+            name="zero_point",
+            shape=(columns, num_groups),
+            initializer="zeros",
+            dtype="uint8",
+            trainable=False,
+        )
+
+        # Per-channel AWQ scales from activation magnitudes
+        self.awq_scales = self.add_weight(
+            name="awq_scales",
+            shape=(rows,),
+            initializer="ones",
+            trainable=False,
+        )
+
+        self.g_idx = self.add_weight(
+            name="g_idx",
+            shape=(rows,),
+            initializer="zeros",
+            dtype="float32",
+            trainable=False,
+        )
+
+    def _awq_call(self, inputs, training=False):
+        """Forward pass for AWQ quantized layer."""
+        if not self.is_awq_calibrated:
+            W = self._kernel
+        else:
+            # Unpack 4-bit weights
+            W = quantizers.unpack_int4(
+                self.quantized_kernel,
+                orig_len=self.awq_unpacked_column_size,
+                axis=0,
+                dtype="uint8",
+            )
+            # Dequantize using scale/zero maps
+            W = dequantize_with_sz_map(
+                W,
+                self.kernel_scale,
+                self.kernel_zero,
+                self.g_idx,
+            )
+            W = ops.transpose(W)
+
+            # Apply AWQ scales by dividing to restore original magnitude
+            # (We multiplied by scales before quantization, so divide to undo)
+            # awq_scales has shape [input_dim], W has shape [input_dim, out_dim]
+            # Expand dims for proper broadcasting.
+            W = ops.divide(W, ops.expand_dims(self.awq_scales, -1))
+
+            W = ops.reshape(W, self.original_kernel_shape)
+
+        y = ops.einsum(self.equation, inputs, W)
+        if self.bias is not None:
+            y = ops.add(y, self.bias)
+        if self.activation is not None:
+            y = self.activation(y)
+        return y
+
     def _int4_build(self, kernel_shape, config=None):
         """Build variables for int4 quantization.
 
@@ -1010,7 +1152,7 @@ class EinsumDense(Layer):
         self.quantization_config = config
 
         kernel_shape = self._kernel.shape
-        if mode in ("int8", "int4", "gptq"):
+        if mode in ("int8", "int4", "gptq", "awq"):
             self._set_quantization_info()
 
         if mode == "int8":
@@ -1058,6 +1200,8 @@ class EinsumDense(Layer):
             policy_name = mode
             if mode == "gptq":
                 policy_name = self.quantization_config.dtype_policy_string()
+            elif mode == "awq":
+                policy_name = self.quantization_config.dtype_policy_string()
             policy = dtype_policies.get(
                 f"{policy_name}_from_{self.dtype_policy.name}"
             )
@@ -1121,7 +1265,7 @@ class EinsumDense(Layer):
                     This is `None` if the layer is not quantized.
         """
         # If not a quantized layer, return the full-precision kernel directly.
-        if self.dtype_policy.quantization_mode in (None, "gptq"):
+        if self.dtype_policy.quantization_mode in (None, "gptq", "awq"):
             return self.kernel, None
 
         # If quantized but LoRA is not enabled, return the original quantized

keras/src/layers/core/embedding.py

@@ -514,7 +514,7 @@ class Embedding(Layer):
                 `embeddings_scale`: The quantization scale for the merged
                     embeddings. This is `None` if the layer is not quantized.
         """
-        if self.dtype_policy.quantization_mode in (None, "gptq"):
+        if self.dtype_policy.quantization_mode in (None, "gptq", "awq"):
             return self.embeddings, None
 
         embeddings_value = self._embeddings

keras/src/layers/core/reversible_embedding.py

@@ -6,6 +6,7 @@ from keras.src import ops
 from keras.src import quantizers
 from keras.src.api_export import keras_export
 from keras.src.backend import KerasTensor
+from keras.src.backend import set_keras_mask
 from keras.src.quantizers.quantization_config import QuantizationConfig
 
 
@@ -117,7 +118,11 @@ class ReversibleEmbedding(layers.Embedding):
 
     def call(self, inputs, reverse=False):
         if not reverse:
-            return super().call(inputs)
+            result = super().call(inputs)
+            mask = super().compute_mask(inputs)
+            if mask is not None:
+                set_keras_mask(result, mask)
+            return result
         else:
             if self.tie_weights:
                 kernel = ops.transpose(ops.convert_to_tensor(self.embeddings))
@@ -135,6 +140,10 @@ class ReversibleEmbedding(layers.Embedding):
                 )
             return logits
 
+    def compute_mask(self, inputs, mask=None):
+        # Disable masking from super class, masking is done directly in call.
+        return None
+
     def compute_output_shape(self, input_shape, reverse=False):
         output_shape = list(input_shape)
         if reverse:

keras/src/layers/layer.py

@@ -1337,6 +1337,8 @@ class Layer(BackendLayer, Operation):
             return self._int4_call(*args, **kwargs)
         elif self.quantization_mode == "gptq":
             return self._gptq_call(*args, **kwargs)
+        elif self.quantization_mode == "awq":
+            return self._awq_call(*args, **kwargs)
         else:
             raise self._quantization_mode_error(self.quantization_mode)
 
@@ -1352,6 +1354,9 @@ class Layer(BackendLayer, Operation):
     def _gptq_call(self, *args, **kwargs):
         raise self._not_implemented_error(self._gptq_call)
 
+    def _awq_call(self, *args, **kwargs):
+        raise self._not_implemented_error(self._awq_call)
+
     def _not_implemented_error(self, attr, msg=None):
         if callable(attr):
             attr_name = attr.__name__

keras/src/layers/preprocessing/feature_space.py

@@ -507,10 +507,14 @@ class FeatureSpace(Layer):
 
     def adapt(self, dataset):
         if not isinstance(dataset, tf.data.Dataset):
-            raise ValueError(
-                "`adapt()` can only be called on a tf.data.Dataset. "
-                f"Received instead: {dataset} (of type {type(dataset)})"
-            )
+            if isinstance(dataset, dict):
+                dataset = tf.data.Dataset.from_tensor_slices(dataset)
+            else:
+                raise ValueError(
+                    "`adapt()` can only be called on a tf.data.Dataset or a "
+                    "dict of arrays/lists. "
+                    f"Received instead: {dataset} (of type {type(dataset)})"
+                )
 
         for name in self._list_adaptable_preprocessors():
             # Call adapt() on each individual adaptable layer.

keras/src/layers/preprocessing/image_preprocessing/aug_mix.py

@@ -316,8 +316,8 @@ class AugMix(BaseImagePreprocessingLayer):
     def get_config(self):
         config = {
             "value_range": self.value_range,
-            "num_chains": self.chain_depth,
-            "chain_depth": self.num_chains,
+            "num_chains": self.num_chains,
+            "chain_depth": self.chain_depth,
             "factor": self.factor,
             "alpha": self.alpha,
             "all_ops": self.all_ops,

keras/src/layers/preprocessing/image_preprocessing/center_crop.py

@@ -183,28 +183,26 @@ class CenterCrop(BaseImagePreprocessingLayer):
 
     def transform_images(self, images, transformation=None, training=True):
         inputs = self.backend.cast(images, self.compute_dtype)
+        inputs_shape = self.backend.shape(inputs)
+
         if self.data_format == "channels_first":
-            init_height = inputs.shape[-2]
-            init_width = inputs.shape[-1]
+            init_height = inputs_shape[-2]
+            init_width = inputs_shape[-1]
         else:
-            init_height = inputs.shape[-3]
-            init_width = inputs.shape[-2]
-
-        if init_height is None or init_width is None:
-            # Dynamic size case. TODO.
-            raise ValueError(
-                "At this time, CenterCrop can only "
-                "process images with a static spatial "
-                f"shape. Received: inputs.shape={inputs.shape}"
-            )
+            init_height = inputs_shape[-3]
+            init_width = inputs_shape[-2]
 
+        # All these operations work both with ints (static sizes) and scalar
+        # tensors (dynamic sizes).
         h_diff = init_height - self.height
         w_diff = init_width - self.width
 
-        h_start = int(h_diff / 2)
-        w_start = int(w_diff / 2)
+        h_start = h_diff // 2
+        w_start = w_diff // 2
 
-        if h_diff >= 0 and w_diff >= 0:
+        if (not isinstance(h_diff, int) or h_diff >= 0) and (
+            not isinstance(w_diff, int) or w_diff >= 0
+        ):
             if len(inputs.shape) == 4:
                 if self.data_format == "channels_first":
                     return inputs[

keras/src/layers/preprocessing/image_preprocessing/random_contrast.py

@@ -92,8 +92,8 @@ class RandomContrast(BaseImagePreprocessingLayer):
 
     def transform_images(self, images, transformation, training=True):
         if training:
-            constrast_factor = transformation["contrast_factor"]
-            outputs = self._adjust_constrast(images, constrast_factor)
+            contrast_factor = transformation["contrast_factor"]
+            outputs = self._adjust_contrast(images, contrast_factor)
             outputs = self.backend.numpy.clip(
                 outputs, self.value_range[0], self.value_range[1]
             )
@@ -117,7 +117,7 @@ class RandomContrast(BaseImagePreprocessingLayer):
     ):
         return segmentation_masks
 
-    def _adjust_constrast(self, inputs, contrast_factor):
+    def _adjust_contrast(self, inputs, contrast_factor):
         if self.data_format == "channels_first":
             height_axis = -2
             width_axis = -1

keras/src/layers/preprocessing/image_preprocessing/resizing.py

@@ -66,6 +66,16 @@ class Resizing(BaseImagePreprocessingLayer):
             `~/.keras/keras.json`. If you never set it, then it will be
             `"channels_last"`.
         **kwargs: Base layer keyword arguments, such as `name` and `dtype`.
+
+    Example:
+
+    ```python
+    (x_train, y_train), _ = keras.datasets.cifar10.load_data()
+    image = x_train[0]
+    resizer = keras.layers.Resizing(128, 128)
+    resized_image = resizer(image)
+    print("original:", image.shape, "resized:", resized_image.shape)
+    ```
     """
 
     _USE_BASE_FACTOR = False

keras/src/losses/losses.py

@@ -73,6 +73,14 @@ class MeanSquaredError(LossFunctionWrapper):
             `"float32"` unless set to different value
             (via `keras.backend.set_floatx()`). If a `keras.DTypePolicy` is
             provided, then the `compute_dtype` will be utilized.
+
+    Examples:
+
+    >>> y_true = keras.ops.array([1.0, 0.0, 1.0])
+    >>> y_pred = keras.ops.array([0.9, 0.1, 0.8])
+    >>> loss = keras.losses.MeanSquaredError()
+    >>> loss(y_true, y_pred)
+    0.02
     """
 
     def __init__(
@@ -114,6 +122,14 @@ class MeanAbsoluteError(LossFunctionWrapper):
             `"float32"` unless set to different value
             (via `keras.backend.set_floatx()`). If a `keras.DTypePolicy` is
             provided, then the `compute_dtype` will be utilized.
+
+    Examples:
+
+    >>> y_true = keras.ops.array([1.0, 0.3, 1.0])
+    >>> y_pred = keras.ops.array([1.9, 0.3, 1.8])
+    >>> loss = keras.losses.MeanAbsoluteError()
+    >>> loss(y_true, y_pred)
+    0.5666667
     """
 
     def __init__(
@@ -155,6 +171,14 @@ class MeanAbsolutePercentageError(LossFunctionWrapper):
             `"float32"` unless set to different value
             (via `keras.backend.set_floatx()`). If a `keras.DTypePolicy` is
             provided, then the `compute_dtype` will be utilized.
+
+    Examples:
+
+    >>> y_true = keras.ops.array([100.0, 200.0, 300.0])
+    >>> y_pred = keras.ops.array([90.0, 210.0, 310.0])
+    >>> loss = keras.losses.MeanAbsolutePercentageError()
+    >>> loss(y_true, y_pred)
+    6.111111
     """
 
     def __init__(