Cherrypick ImageAugmetation related change to R2.9 (keras-team#16454)

* Enable the keras dtensor API in OSS.

PiperOrigin-RevId: 438858608

* Switching learning/brain dependency to OSS compatible test_util

This is one test file failing, due to the monkey patching happens in the dtensor.init(), and I will need to dig more about the root cause (probably due to patching tf.Variable with DVariable, and cause logic difference for instance type checking.)

PiperOrigin-RevId: 439676157

* Update the docstring for keras.dtensor components.

1. Add docstring for LayoutMap.
2. Hide certain methods for keras.dtensor.optimizers.

PiperOrigin-RevId: 442651549

* add no-op label augmentation to RandomContrast

PiperOrigin-RevId: 439938098

* Add augment_label to RandomBrightness

PiperOrigin-RevId: 443122383

* Create augment_label method:
```
  def augment_label(self, label, transformation=None):
    return label
```
on remaining image augmentation layers

PiperOrigin-RevId: 443446864

* Revert "Create augment_label method:"

This reverts commit 1694cb4.

Co-authored-by: Luke Wood <[email protected]>

Cherrypick Keras DTensor related updates into keras 2.9 (keras-team#1…

…6379)

* Enable the keras dtensor API in OSS.

PiperOrigin-RevId: 438858608

* Switching learning/brain dependency to OSS compatible test_util

This is one test file failing, due to the monkey patching happens in the dtensor.init(), and I will need to dig more about the root cause (probably due to patching tf.Variable with DVariable, and cause logic difference for instance type checking.)

PiperOrigin-RevId: 439676157

Change keras to use dtensor public API.

PiperOrigin-RevId: 438605222

Compute LSTM and GRU via cuDNN for RaggedTensors.

This cherry-picked commit fixes a bug preventing computing
LSTM and GRU for RaggedTensors via cuDNN, resulting in
a large speedup (easily 10 times).

A TF2-only test comparing ragged and masked tensor LSTM and GRU
on CPU and GPU is also provided.

Compute LSTM and GRU via cuDNN for RaggedTensors.

This cherry-picked commit fixes a bug preventing computing
LSTM and GRU for RaggedTensors via cuDNN, resulting in
a large speedup (easily 10 times).

A TF2-only test comparing ragged and masked tensor LSTM and GRU
on CPU and GPU is also provided.

Fix RandomCrop layer called on integer input

We had a tf.cond which called tf.image.resize on one path (returning float32) and
cropping the input dtype on the other, which lead to a error. Instead we cast the
inputs to the layer's compute dtype and do our computation with that type.

Fixes keras-team#15681

PiperOrigin-RevId: 417686461

Disable the failing tests due to numpy 1.20 change (keras-team#15552)

The test has been disabled on master branch as well, we need to cherrypick this to unblock the release build.

original cl in cl/400819173.

Disable the failing tests due to numpy 1.20 change (keras-team#15552)

The test has been disabled on master branch as well, we need to cherrypick this to unblock the release build.

original cl in cl/400819173.

Disable the failing tests due to numpy 1.20 change (keras-team#15552)

The test has been disabled on master branch as well, we need to cherrypick this to unblock the release build.

original cl in cl/400819173.

Update Keras Functional model docstring for the new feature of submod…

…el slicing.

PiperOrigin-RevId: 399303217