AugMix layer

keras.layers.AugMix(
    value_range=(0, 255),
    num_chains=3,
    chain_depth=3,
    factor=0.3,
    alpha=1.0,
    all_ops=True,
    interpolation="bilinear",
    seed=None,
    data_format=None,
    **kwargs
)

AugMix aims to produce images with variety while preserving the image
semantics and local statistics. During the augmentation process,
the same augmentation is applied across all images in the batch
in num_chains different ways, with each chain consisting of
chain_depth augmentations.

**Note:** This layer is safe to use inside a [`tf.data`](https://www.tensorflow.org/api_docs/python/tf/data) or `grain` pipeline
(independently of which backend you're using).

References:
    - [AugMix paper](https://arxiv.org/pdf/1912.02781)
    - [Official Code](https://github.com/google-research/augmix)

# Arguments
    value_range: the range of values the incoming images will have.
        Represented as a two number tuple written (low, high).
        This is typically either `(0, 1)` or `(0, 255)` depending
        on how your preprocessing pipeline is set up.
    num_chains: an integer representing the number of different chains to
        be mixed, defaults to 3.
    chain_depth: an integer representing the maximum number of
        transformations to be applied in each chain. The actual number
        of transformations in each chain will be sampled randomly
        from the range `[0, `chain_depth`]`. Defaults to 3.
    factor: The strength of the augmentation as a normalized value
        between 0 and 1. Default is 0.3.
    alpha: a float value used as the probability coefficients for the
        Beta and Dirichlet distributions, defaults to 1.0.
    all_ops: Use all operations (including random_brightness,
        random_color_degeneration, random_contrast and random_sharpness).
        Default is True.
    interpolation: The interpolation method to use for resizing operations.
        Options include `"nearest"`, `"bilinear"`. Default is `"bilinear"`.
    seed: Integer. Used to create a random seed.

# Example

layer = keras.layers.AugMix(bounding_box_format="xyxy")
images = np.random.randint(0, 255, (4, 224, 224, 3), dtype="uint8")

bounding_boxes = {
    "boxes": np.array([
        [[10, 20, 100, 150], [50, 60, 200, 250]],
        [[15, 25, 110, 160], [55, 65, 210, 260]],
        [[20, 30, 120, 170], [60, 70, 220, 270]],
        [[25, 35, 130, 180], [65, 75, 230, 280]],
    ], dtype="float32"),
    "labels": np.array([[0, 1], [1, 2], [2, 3], [0, 3]], dtype="int32")
}

labels = keras.ops.one_hot(
    np.array([0, 1, 2, 3]),
    num_classes=4
)

segmentation_masks = np.random.randint(0, 3, (4, 224, 224, 1), dtype="uint8")

output = layer(
    {
        "images": images,
        "bounding_boxes": bounding_boxes,
        "labels": labels,
        "segmentation_masks": segmentation_masks
    },
    training=True
)