Add support for autocast keyword argument in Layer.add_weight.

This feature was already supported with the `experimental_autocast` argument. This change simply adds an alias for the same argument to have the same API in Keras 2 and Keras 3.

PiperOrigin-RevId: 662606190

tf_keras/engine/base_layer.py

@@ -578,7 +578,8 @@ def add_weight(
  Accepted values are constants defined in the class
  `tf.VariableAggregation`.
  **kwargs: Additional keyword arguments. Accepted values are `getter`,
- `collections`, `experimental_autocast` and `caching_device`.
+ `collections`, `autocast`, `experimental_autocast` and
+ `caching_device`.
 
  Returns:
  The variable created.
@@ -594,6 +595,7 @@ def add_weight(
  # Validate optional keyword arguments.
  for kwarg in kwargs:
  if kwarg not in [
+ "autocast",
  "collections",
  "experimental_autocast",
  "caching_device",
@@ -603,9 +605,13 @@ def add_weight(
  ]:
  raise TypeError("Unknown keyword argument:", kwarg)
  collections_arg = kwargs.pop("collections", None)
- # 'experimental_autocast' can be set to False by the caller to indicate
- # an AutoCastVariable should never be created.
- autocast = kwargs.pop("experimental_autocast", True)
+ # 'autocast' or 'experimental_autocast' can be set to False by the
+ # caller to indicate an AutoCastVariable should never be created.
+ autocast = kwargs.pop("autocast", None)
+ if autocast is None:
+ autocast = kwargs.pop("experimental_autocast", None)
+ if autocast is None:
+ autocast = True
  # See the docstring for tf.Variable about the details for
  # caching_device.
  caching_device = kwargs.pop("caching_device", None)

tf_keras/engine/base_layer_v1.py

@@ -352,7 +352,8 @@ def add_weight(
  Accepted values are constants defined in the class
  `tf.VariableAggregation`.
  **kwargs: Additional keyword arguments. Accepted values are `getter`,
- `collections`, `experimental_autocast` and `caching_device`.
+ `collections`, `autocast`, `experimental_autocast` and
+ `caching_device`.
 
  Returns:
  The created variable. Usually either a `Variable` or
@@ -371,6 +372,7 @@ def add_weight(
  # Validate optional keyword arguments.
  for kwarg in kwargs:
  if kwarg not in [
+ "autocast",
  "getter",
  "collections",
  "experimental_autocast",
@@ -380,9 +382,13 @@ def add_weight(
  has_custom_getter = "getter" in kwargs
  getter = kwargs.pop("getter", base_layer_utils.make_variable)
  collections_arg = kwargs.pop("collections", None)
- # 'experimental_autocast' can be set to False by the caller to indicate
- # an AutoCastVariable should never be created.
- autocast = kwargs.pop("experimental_autocast", True)
+ # 'autocast' or 'experimental_autocast' can be set to False by the
+ # caller to indicate an AutoCastVariable should never be created.
+ autocast = kwargs.pop("autocast", None)
+ if autocast is None:
+ autocast = kwargs.pop("experimental_autocast", None)
+ if autocast is None:
+ autocast = True
  # See the docstring for tf.Variable about the details for
  # caching_device.
  caching_device = kwargs.pop("caching_device", None)

tf_keras/mixed_precision/test_util.py

@@ -214,7 +214,7 @@ def build(self, _):
  (),
  initializer="ones",
  dtype=dtype,
- experimental_autocast=False,
+ autocast=False,
  regularizer=self._regularizer,
  )
  self.built = True