def split_code_into_blocks(lines: List[str], start_index: int, end_index: int, indent: int, backtrace: bool=False) -> List[Tuple[str, int, int]]:
    splits = []
    try:
        target_block_name = re.search(f'^{' ' * (indent - 4)}((class|def)\\s+\\S+)(\\(|\\:)', lines[start_index]).groups()[0]
    except Exception:
        start_context = min(start_index - 10, 0)
        end_context = min(end_index + 10, len(lines))
        raise ValueError(f'Tried to split a class or function. It did not work. Error comes from line {start_index}: \n```\n' + ''.join(lines[start_context:end_context]) + '```\n')
    indent_str = ' ' * indent
    block_without_name_idx = 0
    empty_block_idx = 0
    index = start_index
    if '(' in lines[start_index] and '):' not in lines[start_index] in lines[start_index]:
        while index < end_index:
            if _is_definition_header_ending_line(lines[index]):
                break
            index += 1
    index += 1
    splits.append((target_block_name, start_index, index))
    block_start_index, prev_block_end_index = (index, index)
    while index < end_index:
        block_found = re.search(f'^{indent_str}((class|def)\\s+\\S+)(\\(|\\:)', lines[index])
        if block_found:
            name = block_found.groups()[0]
            block_end_index = find_block_end(lines, index, indent + 4)
            block_start_index = index
            if index > prev_block_end_index and backtrace:
                idx = index - 1
                for idx in range(index - 1, prev_block_end_index - 2, -1):
                    if not (len(lines[idx].strip()) > 0 and lines[idx].startswith(indent_str)):
                        break
                idx += 1
                if idx < index:
                    block_start_index = idx
            if block_start_index > prev_block_end_index:
                if len(''.join(lines[prev_block_end_index:block_start_index]).strip()) == 0:
                    prev_block_name = f'_empty_block_{empty_block_idx}'
                    empty_block_idx += 1
                else:
                    prev_block_name = f'_block_without_name_{block_without_name_idx}'
                    block_without_name_idx += 1
                splits.append((prev_block_name, prev_block_end_index, block_start_index))
            splits.append((name, block_start_index, block_end_index))
            prev_block_end_index = block_end_index
            index = block_end_index - 1
        index += 1
    if index > prev_block_end_index:
        if len(''.join(lines[prev_block_end_index:index]).strip()) == 0:
            prev_block_name = f'_empty_block_{empty_block_idx}'
        else:
            prev_block_name = f'_block_without_name_{block_without_name_idx}'
        splits.append((prev_block_name, prev_block_end_index, index))
    return splits

Split the class/func block starting at `start_index` in a source code (defined by `lines`) into *inner blocks*.

The block's header is included as the first element. The contiguous regions (without empty lines) that are not
inside any inner block are included as blocks. The contiguous regions of empty lines that are not inside any inner
block are also included as (dummy) blocks.

Args:
    lines (`List[str]`):
        The source code, represented by a list of lines.
    start_index (`int`):
        The starting index of the target class/func block.
    end_index (`int`):
        The ending index of the target class/func block.
    indent (`int`):
        The indent of the class/func body.
    backtrace (`bool`, *optional*, defaults to `False`):
        Whether or not to include the lines before the inner class/func block's header (e.g. comments, decorators,
        etc.) until an empty line is encountered.

Returns:
    `List[Tuple[str, int, int]]`: A list of elements with the form `(block_name, start_index, end_index)`.


def build_inputs_with_special_tokens(self, token_ids_0: List[int], token_ids_1: Optional[List[int]]=None) -> List[int]:
    if token_ids_1 is None:
        return [self.cls_token_id] + token_ids_0 + [self.sep_token_id]
    cls = [self.cls_token_id]
    sep = [self.sep_token_id]
    return cls + token_ids_0 + sep + sep + token_ids_1 + sep

Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
adding special tokens. An BARTPho sequence has the following format:

- single sequence: `<s> X </s>`
- pair of sequences: `<s> A </s></s> B </s>`

Args:
    token_ids_0 (`List[int]`):
        List of IDs to which the special tokens will be added.
    token_ids_1 (`List[int]`, *optional*):
        Optional second list of IDs for sequence pairs.

Returns:
    `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.


def forward(self, input):
    output = torch.einsum('eoi,bei->beo', self.weight, input)
    if self.bias is not None:
        raise RuntimeError()
    return output

Args:
    input (`torch.FloatTensor` of shape `(B, n_models, input_dim)`):
        The input to the layer.


class TFCausalLMOutputWithPast(ModelOutput):
    loss: tf.Tensor | None = None
    logits: Optional[tf.Tensor] = None
    past_key_values: List[tf.Tensor] | None = None
    hidden_states: Tuple[tf.Tensor] | None = None
    attentions: Tuple[tf.Tensor] | None = None

Base class for causal language model (or autoregressive) outputs.

Args:
    loss (`tf.Tensor` of shape `(n,)`, *optional*, where n is the number of non-masked labels, returned when `labels` is provided):
        Language modeling loss (for next-token prediction).
    logits (`tf.Tensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
        Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
    past_key_values (`List[tf.Tensor]`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
        List of `tf.Tensor` of length `config.n_layers`, with each tensor of shape `(2, batch_size, num_heads,
        sequence_length, embed_size_per_head)`).

        Contains pre-computed hidden-states (key and values in the attention blocks) that can be used (see
        `past_key_values` input) to speed up sequential decoding.
    hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
        Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape
        `(batch_size, sequence_length, hidden_size)`.

        Hidden-states of the model at the output of each layer plus the initial embedding outputs.
    attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
        Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
        sequence_length)`.

        Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
        heads.


class ASTConfig(PretrainedConfig):
    model_type = 'audio-spectrogram-transformer'

    def __init__(self, hidden_size=768, num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072, hidden_act='gelu', hidden_dropout_prob=0.0, attention_probs_dropout_prob=0.0, initializer_range=0.02, layer_norm_eps=1e-12, patch_size=16, qkv_bias=True, frequency_stride=10, time_stride=10, max_length=1024, num_mel_bins=128, **kwargs):
        super().__init__(**kwargs)
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.intermediate_size = intermediate_size
        self.hidden_act = hidden_act
        self.hidden_dropout_prob = hidden_dropout_prob
        self.attention_probs_dropout_prob = attention_probs_dropout_prob
        self.initializer_range = initializer_range
        self.layer_norm_eps = layer_norm_eps
        self.patch_size = patch_size
        self.qkv_bias = qkv_bias
        self.frequency_stride = frequency_stride
        self.time_stride = time_stride
        self.max_length = max_length
        self.num_mel_bins = num_mel_bins

    def _get_non_default_generation_parameters(self) -> Dict[str, Any]:
        return {}

This is the configuration class to store the configuration of a [`ASTModel`]. It is used to instantiate an AST
model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
defaults will yield a similar configuration to that of the AST
[MIT/ast-finetuned-audioset-10-10-0.4593](https://huggingface.co/MIT/ast-finetuned-audioset-10-10-0.4593)
architecture.

Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
documentation from [`PretrainedConfig`] for more information.

Args:
    hidden_size (`int`, *optional*, defaults to 768):
        Dimensionality of the encoder layers and the pooler layer.
    num_hidden_layers (`int`, *optional*, defaults to 12):
        Number of hidden layers in the Transformer encoder.
    num_attention_heads (`int`, *optional*, defaults to 12):
        Number of attention heads for each attention layer in the Transformer encoder.
    intermediate_size (`int`, *optional*, defaults to 3072):
        Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
    hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
        The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
        `"relu"`, `"selu"` and `"gelu_new"` are supported.
    hidden_dropout_prob (`float`, *optional*, defaults to 0.0):
        The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
    attention_probs_dropout_prob (`float`, *optional*, defaults to 0.0):
        The dropout ratio for the attention probabilities.
    initializer_range (`float`, *optional*, defaults to 0.02):
        The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
    layer_norm_eps (`float`, *optional*, defaults to 1e-12):
        The epsilon used by the layer normalization layers.
    patch_size (`int`, *optional*, defaults to 16):
        The size (resolution) of each patch.
    qkv_bias (`bool`, *optional*, defaults to `True`):
        Whether to add a bias to the queries, keys and values.
    frequency_stride (`int`, *optional*, defaults to 10):
        Frequency stride to use when patchifying the spectrograms.
    time_stride (`int`, *optional*, defaults to 10):
        Temporal stride to use when patchifying the spectrograms.
    max_length (`int`, *optional*, defaults to 1024):
        Temporal dimension of the spectrograms.
    num_mel_bins (`int`, *optional*, defaults to 128):
        Frequency dimension of the spectrograms (number of Mel-frequency bins).

Example:

```python
>>> from transformers import ASTConfig, ASTModel

>>> # Initializing a AST MIT/ast-finetuned-audioset-10-10-0.4593 style configuration
>>> configuration = ASTConfig()

>>> # Initializing a model (with random weights) from the MIT/ast-finetuned-audioset-10-10-0.4593 style configuration
>>> model = ASTModel(configuration)

>>> # Accessing the model configuration
>>> configuration = model.config
```


def sanity_check_tensor_sync(tensor: torch.Tensor, mesh: DeviceMesh, rtol: float=0.0001, atol: float=0.0001, not_sync: bool=False) -> None:
    if not dist.is_initialized() or mesh.size() == 1:
        return
    pg = mesh.get_group()
    if hasattr(tensor, 'to_local'):
        local_tensor = tensor.to_local()
    else:
        local_tensor = tensor
    world_size = dist.get_world_size(pg)
    gathered_tensors = [torch.empty_like(local_tensor) for _ in range(world_size)]
    dist.all_gather(gathered_tensors, local_tensor, group=pg)
    for i in range(1, world_size):
        try:
            torch.testing.assert_close(gathered_tensors[0], gathered_tensors[i], rtol=rtol, atol=atol)
        except AssertionError as e:
            if not_sync:
                continue
            raise e

Verify that a tensor is synchronized (or not synchronized) across all processes in the mesh's process group.
Handles both regular tensors and DTensors.

Args:
    tensor (torch.Tensor): The tensor to check for synchronization (can be DTensor)
    mesh (DeviceMesh): The device mesh containing the process group
    rtol (float): Relative tolerance for comparison
    atol (float): Absolute tolerance for comparison
    not_sync (bool): If True, asserts that tensors are NOT synchronized. If False, asserts they are synchronized.


def get_model_test_files() -> List[str]:
    _ignore_files = ['test_modeling_common', 'test_modeling_encoder_decoder', 'test_modeling_flax_encoder_decoder', 'test_modeling_flax_speech_encoder_decoder', 'test_modeling_marian', 'test_modeling_tf_common', 'test_modeling_tf_encoder_decoder']
    test_files = []
    model_test_root = os.path.join(PATH_TO_TESTS, 'models')
    model_test_dirs = []
    for x in os.listdir(model_test_root):
        x = os.path.join(model_test_root, x)
        if os.path.isdir(x):
            model_test_dirs.append(x)
    for target_dir in [PATH_TO_TESTS] + model_test_dirs:
        for file_or_dir in os.listdir(target_dir):
            path = os.path.join(target_dir, file_or_dir)
            if os.path.isfile(path):
                filename = os.path.split(path)[-1]
                if 'test_modeling' in filename and os.path.splitext(filename)[0] not in _ignore_files:
                    file = os.path.join(*path.split(os.sep)[1:])
                    test_files.append(file)
    return test_files

Get the model test files.

Returns:
    `List[str]`: The list of test files. The returned files will NOT contain the `tests` (i.e. `PATH_TO_TESTS`
    defined in this script). They will be considered as paths relative to `tests`. A caller has to use
    `os.path.join(PATH_TO_TESTS, ...)` to access the files.


def xw_plus_b(x, weights, biases, name=None):
    with ops.name_scope(name, 'xw_plus_b', [x, weights, biases]) as name:
        x = ops.convert_to_tensor(x, name='x')
        weights = ops.convert_to_tensor(weights, name='weights')
        biases = ops.convert_to_tensor(biases, name='biases')
        mm = math_ops.matmul(x, weights)
        return bias_add(mm, biases, name=name)

Computes matmul(x, weights) + biases.

Args:
  x: a 2D tensor.  Dimensions typically: batch, in_units
  weights: a 2D tensor.  Dimensions typically: in_units, out_units
  biases: a 1D tensor.  Dimensions: out_units
  name: A name for the operation (optional).  If not specified
    "xw_plus_b" is used.

Returns:
  A 2-D Tensor computing matmul(x, weights) + biases.
  Dimensions typically: batch, out_units.


def from_function_and_signature(cls, python_function, input_signature, is_pure=False, jit_compile=None):
    function_type, default_values = make_function_type(python_function, input_signature)
    while isinstance(python_function, functools.partial):
        python_function = python_function.func
    name = getattr(python_function, '__name__', 'f')
    return FunctionSpec(function_type, default_values, is_pure=is_pure, jit_compile=jit_compile, name=name)

Creates a FunctionSpec instance given a python function and signature.

Args:
  python_function: a function to inspect
  input_signature: a signature of the function (None, if variable)
  is_pure: if True all input arguments (including variables and constants)
    will be converted to tensors and no variable changes allowed.
  jit_compile: see `tf.function`

Returns:
  instance of FunctionSpec


class JanusCausalLMOutputWithPast(ModelOutput):
    loss: Optional[torch.FloatTensor] = None
    logits: Optional[torch.FloatTensor] = None
    past_key_values: Optional[List[torch.FloatTensor]] = None
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None
    image_hidden_states: Optional[Tuple[torch.FloatTensor]] = None

Base class for Janus causal language model (or autoregressive) outputs.

Args:
    loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
        Language modeling loss (for next-token prediction).
    logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
        Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
    past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
        Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
        `(batch_size, num_heads, sequence_length, embed_size_per_head)`)

        Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
        `past_key_values` input) to speed up sequential decoding.
    hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
        Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
        one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.

        Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
    attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
        Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
        sequence_length)`.

        Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
        heads.
    image_hidden_states (`tuple(torch.FloatTensor)`, *optional*):
        Tuple of `torch.FloatTensor` (one for the output of the image embeddings, `(batch_size, num_images,
        sequence_length, hidden_size)`.

        image_hidden_states of the model produced by the vision encoder, and optionally by the perceiver


def create_model() -> tf.keras.Model:
    return tf.keras.Sequential(pg.oneof([lambda: [tf.keras.layers.Flatten(), tf.keras.layers.Dense(pg.oneof([64, 128]), pg.oneof(['relu', 'sigmoid']))], lambda: [tf.keras.layers.Lambda(lambda x: tf.reshape(x, (-1, 28, 28, 1))), tf.keras.layers.Conv2D(pg.oneof([64, 128]), pg.oneof([(3, 3), (5, 5)]), activation=pg.oneof(['relu', 'sigmoid'])), tf.keras.layers.Flatten()]]) + [tf.keras.layers.Dense(10, activation='softmax')])

Create model for training.

Create a simple tf.keras model for training.

Returns:
  The model to use for training.


def _in_place_subclassed_model_reset(model):
    assert not model._is_graph_network
    version_utils.swap_class(model.__class__, training.Model, training_v1.Model, ops.executing_eagerly_outside_functions())
    attributes_cache = {}
    for name in dir(model):
        if name == 'submodules' or name == '_self_tracked_trackables':
            continue
        try:
            value = getattr(model, name)
        except (AttributeError, ValueError, TypeError):
            continue
        if isinstance(value, Layer):
            attributes_cache[name] = value
            assert value in model.layers
            if hasattr(value, 'layers') and value.layers:
                raise ValueError('We do not support the use of nested layers in `model_to_estimator` at this time. Found nested layer: %s' % value)
        elif isinstance(value, (list, tuple)) and name not in ('layers', '_layers', 'metrics', '_compile_metric_functions', '_output_loss_metrics'):
            if value and all((isinstance(val, Layer) for val in value)):
                raise ValueError('We do not support the use of list-of-layers attributes in subclassed models used with `model_to_estimator` at this time. Found list model: %s' % name)
    layers_to_names = {value: key for key, value in attributes_cache.items()}
    original_layers = list(model._flatten_layers(include_self=False, recursive=False))
    setattr_tracking = model._setattr_tracking
    model._setattr_tracking = False
    model._self_tracked_trackables = []
    for layer in original_layers:
        config = layer.get_config()
        if isinstance(layer, training.Model) and (not layer._is_graph_network):
            raise ValueError('We do not support the use of nested subclassed models in `model_to_estimator` at this time. Found nested model: %s' % layer)
        fresh_layer = layer.__class__.from_config(config)
        name = layers_to_names[layer]
        setattr(model, name, fresh_layer)
        model._self_tracked_trackables.append(fresh_layer)
    if not hasattr(model, '_original_attributes_cache') or model._original_attributes_cache is None:
        if model.built:
            attributes_to_cache = ['inputs', 'outputs', 'total_loss', 'optimizer', 'train_function', 'test_function', 'predict_function', '_training_endpoints', '_collected_trainable_weights', '_feed_inputs', '_feed_input_names', '_feed_input_shapes']
            for name in attributes_to_cache:
                attributes_cache[name] = getattr(model, name)
    model._original_attributes_cache = attributes_cache
    _reset_build_compile_trackers(model)
    model._setattr_tracking = setattr_tracking

Substitute for model cloning that works for subclassed models.

Subclassed models cannot be cloned because their topology is not serializable.
To "instantiate" an identical model in a new TF graph, we reuse the original
model object, but we clear its state.

After calling this function on a model instance, you can use the model
instance as if it were a model clone (in particular you can use it in a new
graph).

This method clears the state of the input model. It is thus destructive.
However the original state can be restored fully by calling
`_in_place_subclassed_model_state_restoration`.

Args:
  model: Instance of a Keras model created via subclassing.

Raises:
  ValueError: In case the model uses a subclassed model as inner layer.


def add_metadata_field(self, field: str, python_type: Type, column_name: Optional[str]=None, convert_fn: Optional[Callable[[Any], Any]]=None, default: Any=None, sql_typecast: Optional[str]=None) -> 'ColumnSpecsBuilder':
    name = column_name or field

    def value_fn(chunk: Chunk) -> Any:
        value = chunk.metadata.get(field, default)
        if value is not None and convert_fn is not None:
            value = convert_fn(value)
        return value
    spec = ColumnSpec(column_name=name, python_type=python_type, value_fn=value_fn, sql_typecast=sql_typecast)
    self._specs.append(spec)
    return self

""Add a :class:`.ColumnSpec` that extracts and converts a field from
chunk metadata.

Args:
    field: Key to extract from chunk metadata
    python_type: Python type for the column (e.g. str, int, float)
    column_name: Name for the column (defaults to metadata field name)
    convert_fn: Optional function to convert the extracted value to
              desired type. If None, value is used as-is
    default: Default value if field is missing from metadata
    sql_typecast: Optional SQL type cast (e.g. "::timestamp")

Returns:
    Self for chaining

Examples:

    Simple string field:
    >>> builder.add_metadata_field("source", str)

    Integer with default:

    >>> builder.add_metadata_field(
    ...     field="count",
    ...     python_type=int,
    ...     column_name="item_count",
    ...     default=0
    ... )

    Float with conversion and default:

    >>> builder.add_metadata_field(
    ...     field="confidence",
    ...     python_type=intfloat,
    ...     convert_fn=lambda x: round(float(x), 2),
    ...     default=0.0
    ... )

    Timestamp with conversion and type cast:

    >>> builder.add_metadata_field(
    ...     field="created_at",
    ...     python_type=intstr,
    ...     convert_fn=lambda ts: ts.replace('T', ' '),
    ...     sql_typecast="::timestamp"
    ... )


class RecallAtPrecision(SensitivitySpecificityBase):

    def __init__(self, precision, num_thresholds=200, class_id=None, name=None, dtype=None):
        if precision < 0 or precision > 1:
            raise ValueError(f'Argument `precision` must be in the range [0, 1]. Received: precision={precision}')
        self.precision = precision
        self.num_thresholds = num_thresholds
        super().__init__(value=precision, num_thresholds=num_thresholds, class_id=class_id, name=name, dtype=dtype)

    def result(self):
        recalls = ops.divide_no_nan(self.true_positives, ops.add(self.true_positives, self.false_negatives))
        precisions = ops.divide_no_nan(self.true_positives, ops.add(self.true_positives, self.false_positives))
        return self._find_max_under_constraint(precisions, recalls, ops.greater_equal)

    def get_config(self):
        config = {'num_thresholds': self.num_thresholds, 'precision': self.precision}
        base_config = super().get_config()
        return {**base_config, **config}

Computes best recall where precision is >= specified value.

For a given score-label-distribution the required precision might not
be achievable, in this case 0.0 is returned as recall.

This metric creates four local variables, `true_positives`,
`true_negatives`, `false_positives` and `false_negatives` that are used to
compute the recall at the given precision. The threshold for the given
precision value is computed and used to evaluate the corresponding recall.

If `sample_weight` is `None`, weights default to 1.
Use `sample_weight` of 0 to mask values.

If `class_id` is specified, we calculate precision by considering only the
entries in the batch for which `class_id` is above the threshold
predictions, and computing the fraction of them for which `class_id` is
indeed a correct label.

Args:
    precision: A scalar value in range `[0, 1]`.
    num_thresholds: (Optional) Defaults to 200. The number of thresholds
        to use for matching the given precision.
    class_id: (Optional) Integer class ID for which we want binary metrics.
        This must be in the half-open interval `[0, num_classes)`, where
        `num_classes` is the last dimension of predictions.
    name: (Optional) string name of the metric instance.
    dtype: (Optional) data type of the metric result.

Example:

>>> m = keras.metrics.RecallAtPrecision(0.8)
>>> m.update_state([0, 0, 1, 1], [0, 0.5, 0.3, 0.9])
>>> m.result()
0.5

>>> m.reset_state()
>>> m.update_state([0, 0, 1, 1], [0, 0.5, 0.3, 0.9],
...                sample_weight=[1, 0, 0, 1])
>>> m.result()
1.0

Usage with `compile()` API:

```python
model.compile(
    optimizer='sgd',
    loss='binary_crossentropy',
    metrics=[keras.metrics.RecallAtPrecision(precision=0.8)])
```


class UniSpeechConfig(PretrainedConfig):
    model_type = 'unispeech'

    def __init__(self, vocab_size=32, hidden_size=768, num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072, hidden_act='gelu', hidden_dropout=0.1, activation_dropout=0.1, attention_dropout=0.1, feat_proj_dropout=0.0, feat_quantizer_dropout=0.0, final_dropout=0.1, layerdrop=0.1, initializer_range=0.02, layer_norm_eps=1e-05, feat_extract_norm='group', feat_extract_activation='gelu', conv_dim=(512, 512, 512, 512, 512, 512, 512), conv_stride=(5, 2, 2, 2, 2, 2, 2), conv_kernel=(10, 3, 3, 3, 3, 2, 2), conv_bias=False, num_conv_pos_embeddings=128, num_conv_pos_embedding_groups=16, do_stable_layer_norm=False, apply_spec_augment=True, mask_time_prob=0.05, mask_time_length=10, mask_time_min_masks=2, mask_feature_prob=0.0, mask_feature_length=10, mask_feature_min_masks=0, num_codevectors_per_group=320, num_codevector_groups=2, contrastive_logits_temperature=0.1, num_negatives=100, codevector_dim=256, proj_codevector_dim=256, diversity_loss_weight=0.1, ctc_loss_reduction='mean', ctc_zero_infinity=False, use_weighted_layer_sum=False, classifier_proj_size=256, num_ctc_classes=80, pad_token_id=0, bos_token_id=1, eos_token_id=2, replace_prob=0.5, **kwargs):
        super().__init__(**kwargs, pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id)
        self.hidden_size = hidden_size
        self.feat_extract_norm = feat_extract_norm
        self.feat_extract_activation = feat_extract_activation
        self.conv_dim = list(conv_dim)
        self.conv_stride = list(conv_stride)
        self.conv_kernel = list(conv_kernel)
        self.conv_bias = conv_bias
        self.num_conv_pos_embeddings = num_conv_pos_embeddings
        self.num_conv_pos_embedding_groups = num_conv_pos_embedding_groups
        self.num_feat_extract_layers = len(self.conv_dim)
        self.num_hidden_layers = num_hidden_layers
        self.intermediate_size = intermediate_size
        self.hidden_act = hidden_act
        self.num_attention_heads = num_attention_heads
        self.hidden_dropout = hidden_dropout
        self.attention_dropout = attention_dropout
        self.activation_dropout = activation_dropout
        self.feat_proj_dropout = feat_proj_dropout
        self.final_dropout = final_dropout
        self.layerdrop = layerdrop
        self.layer_norm_eps = layer_norm_eps
        self.initializer_range = initializer_range
        self.num_ctc_classes = num_ctc_classes
        self.vocab_size = vocab_size
        self.do_stable_layer_norm = do_stable_layer_norm
        self.use_weighted_layer_sum = use_weighted_layer_sum
        self.classifier_proj_size = classifier_proj_size
        if len(self.conv_stride) != self.num_feat_extract_layers or len(self.conv_kernel) != self.num_feat_extract_layers or len(self.conv_dim) != self.num_feat_extract_layers:
            raise ValueError(f'Configuration for convolutional layers is incorrect. It is required that `len(config.conv_dim)` == `len(config.conv_stride)` == `len(config.conv_kernel)`, but is `len(config.conv_dim) = {len(self.conv_dim)}`, `len(config.conv_stride) = {len(self.conv_stride)}`, `len(config.conv_kernel) = {len(self.conv_kernel)}`.')
        self.apply_spec_augment = apply_spec_augment
        self.mask_time_prob = mask_time_prob
        self.mask_time_length = mask_time_length
        self.mask_time_min_masks = mask_time_min_masks
        self.mask_feature_prob = mask_feature_prob
        self.mask_feature_length = mask_feature_length
        self.mask_feature_min_masks = mask_feature_min_masks
        self.num_codevectors_per_group = num_codevectors_per_group
        self.num_codevector_groups = num_codevector_groups
        self.contrastive_logits_temperature = contrastive_logits_temperature
        self.feat_quantizer_dropout = feat_quantizer_dropout
        self.num_negatives = num_negatives
        self.codevector_dim = codevector_dim
        self.proj_codevector_dim = proj_codevector_dim
        self.diversity_loss_weight = diversity_loss_weight
        self.ctc_loss_reduction = ctc_loss_reduction
        self.ctc_zero_infinity = ctc_zero_infinity
        self.replace_prob = replace_prob

    @property
    def inputs_to_logits_ratio(self):
        return functools.reduce(operator.mul, self.conv_stride, 1)

This is the configuration class to store the configuration of a [`UniSpeechModel`]. It is used to instantiate an
UniSpeech model according to the specified arguments, defining the model architecture. Instantiating a
configuration with the defaults will yield a similar configuration to that of the UniSpeech
[microsoft/unispeech-large-1500h-cv](https://huggingface.co/microsoft/unispeech-large-1500h-cv) architecture.

Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
documentation from [`PretrainedConfig`] for more information.

Args:
    vocab_size (`int`, *optional*, defaults to 32):
        Vocabulary size of the UniSpeech model. Defines the number of different tokens that can be represented by
        the `inputs_ids` passed when calling [`UniSpeechModel`]. Vocabulary size of the model. Defines the
        different tokens that can be represented by the *inputs_ids* passed to the forward method of
        [`UniSpeechModel`].
    hidden_size (`int`, *optional*, defaults to 768):
        Dimensionality of the encoder layers and the pooler layer.
    num_hidden_layers (`int`, *optional*, defaults to 12):
        Number of hidden layers in the Transformer encoder.
    num_attention_heads (`int`, *optional*, defaults to 12):
        Number of attention heads for each attention layer in the Transformer encoder.
    intermediate_size (`int`, *optional*, defaults to 3072):
        Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
    hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
        The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
        `"relu"`, `"selu"` and `"gelu_new"` are supported.
    hidden_dropout (`float`, *optional*, defaults to 0.1):
        The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
    activation_dropout (`float`, *optional*, defaults to 0.1):
        The dropout ratio for activations inside the fully connected layer.
    attention_dropout (`float`, *optional*, defaults to 0.1):
        The dropout ratio for the attention probabilities.
    feat_proj_dropout (`float`, *optional*, defaults to 0.0):
        The dropout probability for output of the feature encoder.
    feat_quantizer_dropout (`float`, *optional*, defaults to 0.0):
        The dropout probability for the output of the feature encoder that's used by the quantizer.
    final_dropout (`float`, *optional*, defaults to 0.1):
        The dropout probability for the final projection layer of [`UniSpeechForCTC`].
    layerdrop (`float`, *optional*, defaults to 0.1):
        The LayerDrop probability. See the [LayerDrop paper](see https://huggingface.co/papers/1909.11556) for more
        details.
    initializer_range (`float`, *optional*, defaults to 0.02):
        The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
    layer_norm_eps (`float`, *optional*, defaults to 1e-05):
        The epsilon used by the layer normalization layers.
    feat_extract_norm (`str`, *optional*, defaults to `"group"`):
        The norm to be applied to 1D convolutional layers in feature encoder. One of `"group"` for group
        normalization of only the first 1D convolutional layer or `"layer"` for layer normalization of all 1D
        convolutional layers.
    feat_extract_activation (`str, *optional*, defaults to `"gelu"`):
        The non-linear activation function (function or string) in the 1D convolutional layers of the feature
        extractor. If string, `"gelu"`, `"relu"`, `"selu"` and `"gelu_new"` are supported.
    conv_dim (`Tuple[int]` or `List[int]`, *optional*, defaults to `(512, 512, 512, 512, 512, 512, 512)`):
        A tuple of integers defining the number of input and output channels of each 1D convolutional layer in the
        feature encoder. The length of *conv_dim* defines the number of 1D convolutional layers.
    conv_stride (`Tuple[int]` or `List[int]`, *optional*, defaults to `(5, 2, 2, 2, 2, 2, 2)`):
        A tuple of integers defining the stride of each 1D convolutional layer in the feature encoder. The length
        of *conv_stride* defines the number of convolutional layers and has to match the length of *conv_dim*.
    conv_kernel (`Tuple[int]` or `List[int]`, *optional*, defaults to `(10, 3, 3, 3, 3, 2, 2)`):
        A tuple of integers defining the kernel size of each 1D convolutional layer in the feature encoder. The
        length of *conv_kernel* defines the number of convolutional layers and has to match the length of
        *conv_dim*.
    conv_bias (`bool`, *optional*, defaults to `False`):
        Whether the 1D convolutional layers have a bias.
    num_conv_pos_embeddings (`int`, *optional*, defaults to 128):
        Number of convolutional positional embeddings. Defines the kernel size of 1D convolutional positional
        embeddings layer.
    num_conv_pos_embedding_groups (`int`, *optional*, defaults to 16):
        Number of groups of 1D convolutional positional embeddings layer.
    do_stable_layer_norm (`bool`, *optional*, defaults to `False`):
        Whether to apply *stable* layer norm architecture of the Transformer encoder. `do_stable_layer_norm is
        True` corresponds to applying layer norm before the attention layer, whereas `do_stable_layer_norm is
        False` corresponds to applying layer norm after the attention layer.
    apply_spec_augment (`bool`, *optional*, defaults to `True`):
        Whether to apply *SpecAugment* data augmentation to the outputs of the feature encoder. For reference see
        [SpecAugment: A Simple Data Augmentation Method for Automatic Speech
        Recognition](https://huggingface.co/papers/1904.08779).
    mask_time_prob (`float`, *optional*, defaults to 0.05):
        Percentage (between 0 and 1) of all feature vectors along the time axis which will be masked. The masking
        procedure generates ''mask_time_prob*len(time_axis)/mask_time_length'' independent masks over the axis. If
        reasoning from the probability of each feature vector to be chosen as the start of the vector span to be
        masked, *mask_time_prob* should be `prob_vector_start*mask_time_length`. Note that overlap may decrease the
        actual percentage of masked vectors. This is only relevant if `apply_spec_augment is True`.
    mask_time_length (`int`, *optional*, defaults to 10):
        Length of vector span along the time axis.
    mask_time_min_masks (`int`, *optional*, defaults to 2):
        The minimum number of masks of length `mask_feature_length` generated along the time axis, each time step,
        irrespectively of `mask_feature_prob`. Only relevant if ''mask_time_prob*len(time_axis)/mask_time_length <
        mask_time_min_masks''
    mask_feature_prob (`float`, *optional*, defaults to 0.0):
        Percentage (between 0 and 1) of all feature vectors along the feature axis which will be masked. The
        masking procedure generates ''mask_feature_prob*len(feature_axis)/mask_time_length'' independent masks over
        the axis. If reasoning from the probability of each feature vector to be chosen as the start of the vector
        span to be masked, *mask_feature_prob* should be `prob_vector_start*mask_feature_length`. Note that overlap
        may decrease the actual percentage of masked vectors. This is only relevant if `apply_spec_augment is
        True`.
    mask_feature_length (`int`, *optional*, defaults to 10):
        Length of vector span along the feature axis.
    mask_feature_min_masks (`int`, *optional*, defaults to 0):
        The minimum number of masks of length `mask_feature_length` generated along the feature axis, each time
        step, irrespectively of `mask_feature_prob`. Only relevant if
        ''mask_feature_prob*len(feature_axis)/mask_feature_length < mask_feature_min_masks''
    num_codevectors_per_group (`int`, *optional*, defaults to 320):
        Number of entries in each quantization codebook (group).
    num_codevector_groups (`int`, *optional*, defaults to 2):
        Number of codevector groups for product codevector quantization.
    contrastive_logits_temperature (`float`, *optional*, defaults to 0.1):
        The temperature *kappa* in the contrastive loss.
    num_negatives (`int`, *optional*, defaults to 100):
        Number of negative samples for the contrastive loss.
    codevector_dim (`int`, *optional*, defaults to 256):
        Dimensionality of the quantized feature vectors.
    proj_codevector_dim (`int`, *optional*, defaults to 256):
        Dimensionality of the final projection of both the quantized and the transformer features.
    diversity_loss_weight (`int`, *optional*, defaults to 0.1):
        The weight of the codebook diversity loss component.
    ctc_loss_reduction (`str`, *optional*, defaults to `"mean"`):
        Specifies the reduction to apply to the output of `torch.nn.CTCLoss`. Only relevant when training an
        instance of [`UniSpeechForCTC`].
    ctc_zero_infinity (`bool`, *optional*, defaults to `False`):
        Whether to zero infinite losses and the associated gradients of `torch.nn.CTCLoss`. Infinite losses mainly
        occur when the inputs are too short to be aligned to the targets. Only relevant when training an instance
        of [`UniSpeechForCTC`].
    use_weighted_layer_sum (`bool`, *optional*, defaults to `False`):
        Whether to use a weighted average of layer outputs with learned weights. Only relevant when using an
        instance of [`UniSpeechForSequenceClassification`].
    classifier_proj_size (`int`, *optional*, defaults to 256):
        Dimensionality of the projection before token mean-pooling for classification.
    num_ctc_classes (`int`, *optional*, defaults to 80):
        Specifies the number of classes (phoneme tokens and blank token) for phoneme-level CTC loss. Only relevant
        when using an instance of [`UniSpeechForPreTraining`].
    pad_token_id (`int`, *optional*, defaults to 0):
        The id of the padding token.
    bos_token_id (`int`, *optional*, defaults to 1):
        The id of the "beginning-of-sequence" token.
    eos_token_id (`int`, *optional*, defaults to 2):
        The id of the "end-of-sequence" token.
    replace_prob (`float`, *optional*, defaults to 0.5):
        Probability that transformer feature is replaced by quantized feature for pretraining.

Example:

```python
>>> from transformers import UniSpeechConfig, UniSpeechModel

>>> # Initializing a UniSpeech facebook/unispeech-base-960h style configuration
>>> configuration = UniSpeechConfig()

>>> # Initializing a model (with random weights) from the facebook/unispeech-base-960h style configuration
>>> model = UniSpeechModel(configuration)

>>> # Accessing the model configuration
>>> configuration = model.config
```


def _partitioner(shape, dtype):
    if not isinstance(shape, tensor_shape.TensorShape):
        raise ValueError(f'shape is not a TensorShape: {shape}')
    if not shape.is_fully_defined():
        raise ValueError(f'shape is not fully defined: {shape}')
    dtype = dtypes.as_dtype(dtype)
    if dtype.base_dtype == dtypes.string:
        element_size = bytes_per_string_element
    else:
        element_size = dtype.size
    partitions = [1] * shape.ndims
    bytes_per_slice = 1.0 * (shape.num_elements() / shape.dims[axis].value) * element_size
    slices_per_shard = max(1, math.floor(max_shard_bytes / bytes_per_slice))
    axis_shards = int(math.ceil(1.0 * shape.dims[axis].value / slices_per_shard))
    if max_shards:
        axis_shards = min(max_shards, axis_shards)
    partitions[axis] = axis_shards
    return partitions

Partitioner that partitions shards to have max_shard_bytes total size.

Args:
  shape: A `TensorShape`.
  dtype: A `DType`.

Returns:
  A tuple representing how much to slice each axis in shape.

Raises:
  ValueError: If shape is not a fully defined `TensorShape` or dtype is not
    a `DType`.


def batch_shape_tensor(self, name='batch_shape_tensor'):
    with self._name_scope(name):
        if self.batch_shape.is_fully_defined():
            return ops.convert_to_tensor(self.batch_shape.as_list(), dtype=dtypes.int32, name='batch_shape')
        return self._batch_shape_tensor()

Shape of a single sample from a single event index as a 1-D `Tensor`.

The batch dimensions are indexes into independent, non-identical
parameterizations of this distribution.

Args:
  name: name to give to the op

Returns:
  batch_shape: `Tensor`.


def __init__(self, config, auth, trix_id, feed_name, parse=True, spreadsheet=None, timezone=None):
    self.config = config
    self.auth = auth
    self.trix_id = trix_id
    self.trix_range = 'A1:AZ'
    self.feed_name = feed_name
    self._parse = parse
    self._timezone = timezone or 'America/New_York'
    if spreadsheet:
        self.spreadsheet = spreadsheet
    else:
        self.spreadsheet = sheets_get(self.config, self.auth, self.trix_id)
    self.raw_feed = self._get_feed()
    self.feed = self._feed_to_dict(parse=self._parse)

Initializes the feed with parameters.

Args:
  auth: The authentication scheme to use based on the json configuration
    file.
  trix_id: Unique identifier of the Google Sheet that represents the
    Bulkdozer feed.
  feed_name: The name of the feed to initialize.
  spreadsheet: Optional, the spreadsheet object representing the Bulkdozer
    feed spreadsheet, useful to limit calls to the sheets API and allow
    multiple Feed objects to use the same spreadsheet instance. This is used
    to determine which tabs exist in the feed so the correct one can be
    selected for the entity this Feed object represents.


class OptimizerV2(trackable.Trackable):
    _HAS_AGGREGATE_GRAD = False

    def __init__(self, name, gradient_aggregator=None, gradient_transformers=None, **kwargs):
        """Create a new Optimizer.

    This must be called by the constructors of subclasses.
    Note that Optimizer instances should not bind to a single graph,
    and so shouldn't keep Tensors as member variables. Generally
    you should be able to use the _set_hyper()/state.get_hyper()
    facility instead.

    This class is stateful and thread-compatible.

    Example of custom gradient transformations:

    ```python
    def my_gradient_transformer(grads_and_vars):
      # Simple example, double the gradients.
      return [(2. * g, v) for g, v in grads_and_vars]

    optimizer = tf.keras.optimizers.SGD(
        1e-3, gradient_transformers=[my_gradient_transformer])
    ```

    Args:
      name: String. The name to use for momentum accumulator weights created
        by the optimizer.
      gradient_aggregator: The function to use to aggregate gradients across
        devices (when using `tf.distribute.Strategy`). If `None`, defaults to
        summing the gradients across devices. The function should accept and
        return a list of `(gradient, variable)` tuples.
      gradient_transformers: Optional. List of functions to use to transform
        gradients before applying updates to Variables. The functions are
        applied after `gradient_aggregator`. The functions should accept and
        return a list of `(gradient, variable)` tuples.
      **kwargs: keyword arguments. Allowed arguments are `clipvalue`,
        `clipnorm`, `global_clipnorm`.
        If `clipvalue` (float) is set, the gradient of each weight
        is clipped to be no higher than this value.
        If `clipnorm` (float) is set, the gradient of each weight
        is individually clipped so that its norm is no higher than this value.
        If `global_clipnorm` (float) is set the gradient of all weights is
        clipped so that their global norm is no higher than this value.

    Raises:
      ValueError: in case of any invalid argument.
    """
        allowed_kwargs = {'clipnorm', 'clipvalue', 'lr', 'decay', 'global_clipnorm'}
        for k in kwargs:
            if k not in allowed_kwargs:
                raise TypeError('Unexpected keyword argument passed to optimizer: ' + str(k))
            if kwargs[k] is not None and kwargs[k] < 0:
                raise ValueError('Expected {} >= 0, received: {}'.format(k, kwargs[k]))
            if k == 'lr':
                warnings.warn('The `lr` argument is deprecated, use `learning_rate` instead.')
        self._use_locking = True
        self._init_set_name(name)
        self._hyper = {}
        self._slots = {}
        self._slot_names = []
        self._weights = []
        self._iterations = None
        self._deferred_slot_restorations = {}
        decay = kwargs.pop('decay', 0.0)
        if decay < 0.0:
            raise ValueError('decay cannot be less than 0: {}'.format(decay))
        self._initial_decay = decay
        self._hypers_created = False
        if distribute_lib.has_strategy():
            self._distribution_strategy = distribute_lib.get_strategy()
        else:
            self._distribution_strategy = None
        if gradient_aggregator is None:
            gradient_aggregator = optimizer_utils.all_reduce_sum_gradients
        self.gradient_aggregator = gradient_aggregator
        if gradient_transformers is None:
            gradient_transformers = []
        self.gradient_transformers = gradient_transformers
        self.clipnorm = kwargs.pop('clipnorm', None)
        self.global_clipnorm = kwargs.pop('global_clipnorm', None)
        if self.clipnorm is not None and self.global_clipnorm is not None:
            raise ValueError('Cannot accept both `clipnorm` and `global_clipnorm`, passed `clipnorm` {}, `global_clipnorm` {}'.format(self.clipnorm, self.global_clipnorm))
        self.clipvalue = kwargs.pop('clipvalue', None)

    @property
    def clipnorm(self):
        """`float` or `None`. If set, clips gradients to a maximum norm."""
        return self._clipnorm

    @property
    def global_clipnorm(self):
        """`float` or `None`. If set, clips gradients to a maximum norm."""
        return self._global_clipnorm

    @clipnorm.setter
    def clipnorm(self, val):
        if val is not None and self.gradient_transformers:
            raise ValueError('`clipnorm` cannot be set when `gradient_transformers` is set. Instead, use the `gradient_transformers` to specify clipping and other transformations.')
        self._clipnorm = val
        self._clipnorm_fn = optimizer_utils.make_gradient_clipnorm_fn(self._clipnorm)

    @global_clipnorm.setter
    def global_clipnorm(self, val):
        if val is not None and self.gradient_transformers:
            raise ValueError('`clipnorm` cannot be set when `gradient_transformers` is set. Instead, use the `gradient_transformers` to specify clipping and other transformations.')
        self._global_clipnorm = val
        self._global_clipnorm_fn = optimizer_utils.make_global_gradient_clipnorm_fn(self._global_clipnorm)

    @property
    def clipvalue(self):
        """`float` or `None`. If set, clips gradients to a maximum value."""
        return self._clipvalue

    @clipvalue.setter
    def clipvalue(self, val):
        if val is not None and self.gradient_transformers:
            raise ValueError('`clipvalue` cannot be set when `gradient_transformers` is set. Instead, use the `gradient_transformers` to specify clipping and other transformations.')
        self._clipvalue = val
        self._clipvalue_fn = optimizer_utils.make_gradient_clipvalue_fn(self._clipvalue)

    def _transform_loss(self, loss):
        """Called in `.minimize` to transform loss before computing gradients."""
        return loss

    def _get_gradients(self, tape, loss, var_list, grad_loss=None):
        """Called in `minimize` to compute gradients from loss."""
        grads = tape.gradient(loss, var_list, grad_loss)
        return list(zip(grads, var_list))

    def _transform_unaggregated_gradients(self, grads_and_vars):
        """Called in `apply_gradients` before gradient aggregation."""
        return grads_and_vars

    def _aggregate_gradients(self, grads_and_vars):
        """Called in `apply_gradients` to aggregate gradients across devices.

    Note that user subclasses may override this, so the interface should not be
    changed.

    Args:
      grads_and_vars: List of (gradient, variable) pairs.

    Returns:
      A list of (aggregated_gradient, variable) pairs. By default, this calls
      `self.gradient_aggregator`.
    """
        return self.gradient_aggregator(grads_and_vars)

    def _transform_gradients(self, grads_and_vars):
        """Called in `apply_gradients` after aggregation."""
        if self._clipvalue is not None:
            grads_and_vars = self._clipvalue_fn(grads_and_vars)
        if self._clipnorm is not None:
            grads_and_vars = self._clipnorm_fn(grads_and_vars)
        if self._global_clipnorm is not None:
            grads_and_vars = self._global_clipnorm_fn(grads_and_vars)
        for fn in self.gradient_transformers:
            grads_and_vars = fn(grads_and_vars)
        return grads_and_vars

    def minimize(self, loss, var_list, grad_loss=None, name=None, tape=None):
        """Minimize `loss` by updating `var_list`.

    This method simply computes gradient using `tf.GradientTape` and calls
    `apply_gradients()`. If you want to process the gradient before applying
    then call `tf.GradientTape` and `apply_gradients()` explicitly instead
    of using this function.

    Args:
      loss: `Tensor` or callable. If a callable, `loss` should take no arguments
        and return the value to minimize. If a `Tensor`, the `tape` argument
        must be passed.
      var_list: list or tuple of `Variable` objects to update to minimize
        `loss`, or a callable returning the list or tuple of `Variable` objects.
        Use callable when the variable list would otherwise be incomplete before
        `minimize` since the variables are created at the first time `loss` is
        called.
      grad_loss: (Optional). A `Tensor` holding the gradient computed for
        `loss`.
      name: (Optional) str. Name for the returned operation.
      tape: (Optional) `tf.GradientTape`. If `loss` is provided as a `Tensor`,
        the tape that computed the `loss` must be provided.

    Returns:
      An `Operation` that updates the variables in `var_list`. The `iterations`
      will be automatically increased by 1.

    Raises:
      ValueError: If some of the variables are not `Variable` objects.

    """
        grads_and_vars = self._compute_gradients(loss, var_list=var_list, grad_loss=grad_loss, tape=tape)
        return self.apply_gradients(grads_and_vars, name=name)

    def _compute_gradients(self, loss, var_list, grad_loss=None, tape=None):
        """Compute gradients of `loss` for the variables in `var_list`.

    This is the first part of `minimize()`.  It returns a list
    of (gradient, variable) pairs where "gradient" is the gradient
    for "variable".  Note that "gradient" can be a `Tensor`, an
    `IndexedSlices`, or `None` if there is no gradient for the
    given variable.

    Args:
      loss: `Tensor` or callable. If a callable, `loss` should take no
        arguments and return the value to minimize. If a `Tensor`, the `tape`
        argument must be passed.
      var_list: list or tuple of `Variable` objects to update to minimize
        `loss`, or a callable returning the list or tuple of `Variable` objects.
        Use callable when the variable list would otherwise be incomplete before
        `minimize` and the variables are created at the first time when `loss`
        is called.
      grad_loss: Optional. A `Tensor` holding the gradient computed for `loss`.
      tape: (Optional) `tf.GradientTape`. If `loss` is provided as a `Tensor`,
        the tape that computed the `loss` must be provided.

    Returns:
      A list of (gradient, variable) pairs. Variable is always present, but
      gradient can be `None`.

    Raises:
      TypeError: If `var_list` contains anything else than `Variable` objects.
      ValueError: If some arguments are invalid, or var_list is None.
    """
        if not callable(loss) and tape is None:
            raise ValueError('`tape` is required when a `Tensor` loss is passed.')
        tape = tape if tape is not None else backprop.GradientTape()
        if callable(loss):
            with tape:
                if not callable(var_list):
                    tape.watch(var_list)
                loss = loss()
                if callable(var_list):
                    var_list = var_list()
        with tape:
            loss = self._transform_loss(loss)
        var_list = nest.flatten(var_list)
        with ops.name_scope_v2(self._name + '/gradients'):
            grads_and_vars = self._get_gradients(tape, loss, var_list, grad_loss)
        self._assert_valid_dtypes([v for g, v in grads_and_vars if g is not None and v.dtype != dtypes.resource])
        return grads_and_vars

    def apply_gradients(self, grads_and_vars, name=None, experimental_aggregate_gradients=True):
        """Apply gradients to variables.

    This is the second part of `minimize()`. It returns an `Operation` that
    applies gradients.

    The method sums gradients from all replicas in the presence of
    `tf.distribute.Strategy` by default. You can aggregate gradients yourself by
    passing `experimental_aggregate_gradients=False`.

    Example:

    ```python
    grads = tape.gradient(loss, vars)
    grads = tf.distribute.get_replica_context().all_reduce('sum', grads)
    # Processing aggregated gradients.
    optimizer.apply_gradients(zip(grads, vars),
        experimental_aggregate_gradients=False)

    ```

    Args:
      grads_and_vars: List of (gradient, variable) pairs.
      name: Optional name for the returned operation. Default to the name passed
        to the `Optimizer` constructor.
      experimental_aggregate_gradients: Whether to sum gradients from different
        replicas in the presence of `tf.distribute.Strategy`. If False, it's
        user responsibility to aggregate the gradients. Default to True.

    Returns:
      An `Operation` that applies the specified gradients. The `iterations`
      will be automatically increased by 1.

    Raises:
      TypeError: If `grads_and_vars` is malformed.
      ValueError: If none of the variables have gradients.
      RuntimeError: If called in a cross-replica context.
    """
        grads_and_vars = optimizer_utils.filter_empty_gradients(grads_and_vars)
        var_list = [v for _, v in grads_and_vars]
        with ops.name_scope_v2(self._name):
            with ops.init_scope():
                self._create_all_weights(var_list)
            if not grads_and_vars:
                return control_flow_ops.no_op()
            if distribute_lib.in_cross_replica_context():
                raise RuntimeError('`apply_gradients() cannot be called in cross-replica context. Use `tf.distribute.Strategy.run` to enter replica context.')
            strategy = distribute_lib.get_strategy()
            if not experimental_aggregate_gradients and strategy and isinstance(strategy, (parameter_server_strategy.ParameterServerStrategyV1, parameter_server_strategy_v2.ParameterServerStrategyV2, central_storage_strategy.CentralStorageStrategy, central_storage_strategy.CentralStorageStrategyV1)):
                raise NotImplementedError('`experimental_aggregate_gradients=False is not supported for ParameterServerStrategy and CentralStorageStrategy')
            apply_state = self._prepare(var_list)
            if experimental_aggregate_gradients:
                grads_and_vars = self._transform_unaggregated_gradients(grads_and_vars)
                grads_and_vars = self._aggregate_gradients(grads_and_vars)
            grads_and_vars = self._transform_gradients(grads_and_vars)
            if optimizer_utils.strategy_supports_no_merge_call():
                return self._distributed_apply(strategy, grads_and_vars, name, apply_state)
            else:
                return distribute_lib.get_replica_context().merge_call(functools.partial(self._distributed_apply, apply_state=apply_state), args=(grads_and_vars,), kwargs={'name': name})

    def _distributed_apply(self, distribution, grads_and_vars, name, apply_state):
        """`apply_gradients` using a `DistributionStrategy`."""

        def apply_grad_to_update_var(var, grad):
            """Apply gradient to variable."""
            if isinstance(var, tensor.Tensor):
                raise NotImplementedError('Trying to update a Tensor ', var)
            apply_kwargs = {}
            if isinstance(grad, indexed_slices.IndexedSlices):
                if var.constraint is not None:
                    raise RuntimeError('Cannot use a constraint function on a sparse variable.')
                if 'apply_state' in self._sparse_apply_args:
                    apply_kwargs['apply_state'] = apply_state
                return self._resource_apply_sparse_duplicate_indices(grad.values, var, grad.indices, **apply_kwargs)
            if 'apply_state' in self._dense_apply_args:
                apply_kwargs['apply_state'] = apply_state
            update_op = self._resource_apply_dense(grad, var, **apply_kwargs)
            if var.constraint is not None:
                with ops.control_dependencies([update_op]):
                    return var.assign(var.constraint(var))
            else:
                return update_op
        eagerly_outside_functions = ops.executing_eagerly_outside_functions()
        update_ops = []
        with name_scope_only_in_function_or_graph(name or self._name):
            for grad, var in grads_and_vars:
                with distribution.extended.colocate_vars_with(var):
                    with name_scope_only_in_function_or_graph('update' if eagerly_outside_functions else 'update_' + var.op.name):
                        update_op = distribution.extended.update(var, apply_grad_to_update_var, args=(grad,), group=False)
                        if distribute_lib.in_cross_replica_context():
                            update_ops.extend(update_op)
                        else:
                            update_ops.append(update_op)
            any_symbolic = any((isinstance(i, ops.Operation) or tf_utils.is_symbolic_tensor(i) for i in update_ops))
            if not context.executing_eagerly() or any_symbolic:
                with backend._current_graph(update_ops).as_default():
                    with ops.control_dependencies([control_flow_ops.group(update_ops)]):
                        return self._iterations.assign_add(1, read_value=False)
            return self._iterations.assign_add(1)

    def get_gradients(self, loss, params):
        """Returns gradients of `loss` with respect to `params`.

    Should be used only in legacy v1 graph mode.

    Args:
      loss: Loss tensor.
      params: List of variables.

    Returns:
      List of gradient tensors.

    Raises:
      ValueError: In case any gradient cannot be computed (e.g. if gradient
        function not implemented).
    """
        params = nest.flatten(params)
        with backend.get_graph().as_default(), backend.name_scope(self._name + '/gradients'):
            grads = gradients.gradients(loss, params)
            for grad, param in zip(grads, params):
                if grad is None:
                    raise ValueError('Variable {} has `None` for gradient. Please make sure that all of your ops have a gradient defined (i.e. are differentiable). Common ops without gradient: K.argmax, K.round, K.eval.'.format(param))
        return grads

    def get_updates(self, loss, params):
        grads = self.get_gradients(loss, params)
        grads_and_vars = list(zip(grads, params))
        self._assert_valid_dtypes([v for g, v in grads_and_vars if g is not None and v.dtype != dtypes.resource])
        return [self.apply_gradients(grads_and_vars)]

    def _set_hyper(self, name, value):
        """set hyper `name` to value. value can be callable, tensor, numeric."""
        if isinstance(value, trackable.Trackable):
            self._track_trackable(value, name, overwrite=True)
        if name not in self._hyper:
            self._hyper[name] = value
        else:
            prev_value = self._hyper[name]
            if callable(prev_value) or isinstance(prev_value, (tensor.Tensor, int, float, learning_rate_schedule.LearningRateSchedule)) or isinstance(value, learning_rate_schedule.LearningRateSchedule):
                self._hyper[name] = value
            else:
                backend.set_value(self._hyper[name], value)

    def _get_hyper(self, name, dtype=None):
        if not self._hypers_created:
            self._create_hypers()
        value = self._hyper[name]
        if isinstance(value, learning_rate_schedule.LearningRateSchedule):
            return value
        if callable(value):
            value = value()
        if dtype:
            return math_ops.cast(value, dtype)
        else:
            return value

    def _create_slots(self, var_list):
        pass

    def _create_all_weights(self, var_list):
        """Creates all weights, including iterations, hyperparameters and slot vars.

    This will add newly created variables to `optimizer.weights`.

    New variables are only created when this method is called the first time, or
    when called with different variables in the var_list.

    Args:
      var_list: list or tuple of `Variable` objects that will be minimized
        using this optimizer.
    """
        _ = self.iterations
        self._create_hypers()
        self._create_slots(var_list)

    def __getattribute__(self, name):
        """Overridden to support hyperparameter access."""
        try:
            return super(OptimizerV2, self).__getattribute__(name)
        except AttributeError as e:
            if name == '_hyper':
                raise e
            if name == 'lr':
                name = 'learning_rate'
            if name in self._hyper:
                return self._get_hyper(name)
            raise e

    def __dir__(self):
        result = set(super(OptimizerV2, self).__dir__())
        if '_hyper' in result:
            result |= self._hyper.keys()
            if 'learning_rate' in self._hyper.keys():
                result.add('lr')
        return list(result)

    def __setattr__(self, name, value):
        """Override setattr to support dynamic hyperparameter setting."""
        if name == 'lr':
            name = 'learning_rate'
        if hasattr(self, '_hyper') and name in self._hyper:
            self._set_hyper(name, value)
        else:
            super(OptimizerV2, self).__setattr__(name, value)

    def get_slot_names(self):
        """A list of names for this optimizer's slots."""
        return self._slot_names

    def add_slot(self, var, slot_name, initializer='zeros', shape=None):
        """Add a new slot variable for `var`.

    A slot variable is an additional variable associated with `var` to train.
    It is allocated and managed by optimizers, e.g. `Adam`.

    Args:
      var: a `Variable` object.
      slot_name: name of the slot variable.
      initializer: initializer of the slot variable
      shape: (Optional) shape of the slot variable. If not set, it will default
      to the shape of `var`.

    Returns:
      A slot variable.
    """
        if slot_name not in self._slot_names:
            self._slot_names.append(slot_name)
        var_key = _var_key(var)
        slot_dict = self._slots.setdefault(var_key, {})
        weight = slot_dict.get(slot_name, None)
        if weight is None:
            if isinstance(initializer, str) or callable(initializer):
                initializer = initializers.get(initializer)
                if isinstance(initializer, trackable.CheckpointInitialValueCallable) or shape is not None:
                    slot_shape = shape
                else:
                    slot_shape = var.shape
                initial_value = functools.partial(initializer, shape=slot_shape, dtype=var.dtype)
            else:
                initial_value = initializer
            with self._distribution_strategy_scope():
                strategy = distribute_lib.get_strategy()
                if not strategy.extended.variable_created_in_scope(var):
                    raise ValueError("Trying to create optimizer slot variable under the scope for tf.distribute.Strategy ({}), which is different from the scope used for the original variable ({}). Make sure the slot variables are created under the same strategy scope. This may happen if you're restoring from a checkpoint outside the scope".format(strategy, var))
                with strategy.extended.colocate_vars_with(var):
                    weight = tf_variables.Variable(name='%s/%s' % (var._shared_name, slot_name), dtype=var.dtype, trainable=False, initial_value=initial_value)
            backend.track_variable(weight)
            slot_dict[slot_name] = weight
            self._restore_slot_variable(slot_name=slot_name, variable=var, slot_variable=weight)
            self._weights.append(weight)
        return weight

    def get_slot(self, var, slot_name):
        var_key = _var_key(var)
        slot_dict = self._slots[var_key]
        return slot_dict[slot_name]

    def _prepare(self, var_list):
        keys = set()
        for var in var_list:
            if isinstance(var, ds_values.DistributedValues):
                var_devices = var._devices
            else:
                var_devices = [var.device]
            var_dtype = var.dtype.base_dtype
            for var_device in var_devices:
                keys.add((var_device, var_dtype))
        apply_state = {}
        for var_device, var_dtype in keys:
            apply_state[var_device, var_dtype] = {}
            with ops.device(var_device):
                self._prepare_local(var_device, var_dtype, apply_state)
        return apply_state

    def _prepare_local(self, var_device, var_dtype, apply_state):
        if 'learning_rate' in self._hyper:
            lr_t = array_ops.identity(self._decayed_lr(var_dtype))
            apply_state[var_device, var_dtype]['lr_t'] = lr_t

    def _fallback_apply_state(self, var_device, var_dtype):
        """Compatibility for subclasses that don't pass apply_state through."""
        apply_state = {(var_device, var_dtype): {}}
        self._prepare_local(var_device, var_dtype, apply_state)
        return apply_state[var_device, var_dtype]

    def _create_hypers(self):
        if self._hypers_created:
            return
        with self._distribution_strategy_scope():
            for name, value in sorted(self._hyper.items()):
                if isinstance(value, (tensor.Tensor, tf_variables.Variable)) or callable(value):
                    continue
                else:
                    self._hyper[name] = self.add_weight(name, shape=[], trainable=False, initializer=value, aggregation=tf_variables.VariableAggregation.ONLY_FIRST_REPLICA)
        self._hypers_created = True

    @property
    def iterations(self):
        """Variable. The number of training steps this Optimizer has run."""
        if self._iterations is None:
            with self._distribution_strategy_scope():
                self._iterations = self.add_weight('iter', shape=[], dtype=dtypes.int64, trainable=False, aggregation=tf_variables.VariableAggregation.ONLY_FIRST_REPLICA)
            self._weights.append(self._iterations)
        return self._iterations

    @iterations.setter
    def iterations(self, variable):
        if self._iterations is not None:
            raise RuntimeError('Cannot set `iterations` to a new Variable after the Optimizer weights have been created')
        self._iterations = variable
        self._weights.append(self._iterations)

    def _decayed_lr(self, var_dtype):
        """Get decayed learning rate as a Tensor with dtype=var_dtype."""
        lr_t = self._get_hyper('learning_rate', var_dtype)
        if isinstance(lr_t, learning_rate_schedule.LearningRateSchedule):
            local_step = math_ops.cast(self.iterations, var_dtype)
            lr_t = math_ops.cast(lr_t(local_step), var_dtype)
        if self._initial_decay > 0.0:
            local_step = math_ops.cast(self.iterations, var_dtype)
            decay_t = math_ops.cast(self._initial_decay, var_dtype)
            lr_t = lr_t / (1.0 + decay_t * local_step)
        return lr_t

    @abc.abstractmethod
    def get_config(self):
        """Returns the config of the optimizer.

    An optimizer config is a Python dictionary (serializable)
    containing the configuration of an optimizer.
    The same optimizer can be reinstantiated later
    (without any saved state) from this configuration.

    Returns:
        Python dictionary.
    """
        config = {'name': self._name}
        if self.clipnorm is not None:
            config['clipnorm'] = self.clipnorm
        if self.clipvalue is not None:
            config['clipvalue'] = self.clipvalue
        if self.global_clipnorm is not None:
            config['global_clipnorm'] = self.global_clipnorm
        return config

    @classmethod
    def from_config(cls, config, custom_objects=None):
        """Creates an optimizer from its config.

    This method is the reverse of `get_config`,
    capable of instantiating the same optimizer from the config
    dictionary.

    Args:
        config: A Python dictionary, typically the output of get_config.
        custom_objects: A Python dictionary mapping names to additional Python
          objects used to create this optimizer, such as a function used for a
          hyperparameter.

    Returns:
        An optimizer instance.
    """
        if 'lr' in config:
            config['learning_rate'] = config.pop('lr')
        if 'learning_rate' in config:
            if isinstance(config['learning_rate'], dict):
                config['learning_rate'] = learning_rate_schedule.deserialize(config['learning_rate'], custom_objects=custom_objects)
        return cls(**config)

    def _serialize_hyperparameter(self, hyperparameter_name):
        """Serialize a hyperparameter that can be a float, callable, or Tensor."""
        value = self._hyper[hyperparameter_name]
        if isinstance(value, learning_rate_schedule.LearningRateSchedule):
            return learning_rate_schedule.serialize(value)
        if callable(value):
            return value()
        if tensor_util.is_tf_type(value):
            return backend.get_value(value)
        return value

    def variables(self):
        """Returns variables of this Optimizer based on the order created."""
        return self._weights

    @property
    def weights(self):
        """Returns variables of this Optimizer based on the order created."""
        return self._weights

    def get_weights(self):
        """Returns the current weights of the optimizer.

    The weights of an optimizer are its state (ie, variables).
    This function returns the weight values associated with this
    optimizer as a list of Numpy arrays. The first value is always the
    iterations count of the optimizer, followed by the optimizer's state
    variables in the order they were created. The returned list can in turn
    be used to load state into similarly parameterized optimizers.

    For example, the RMSprop optimizer for this simple model returns a list of
    three values-- the iteration count, followed by the root-mean-square value
    of the kernel and bias of the single Dense layer:

    >>> opt = tf.keras.optimizers.RMSprop()
    >>> m = tf.keras.models.Sequential([tf.keras.layers.Dense(10)])
    >>> m.compile(opt, loss='mse')
    >>> data = np.arange(100).reshape(5, 20)
    >>> labels = np.zeros(5)
    >>> results = m.fit(data, labels)  # Training.
    >>> len(opt.get_weights())
    3

    Returns:
        Weights values as a list of numpy arrays.
    """
        params = self.weights
        return backend.batch_get_value(params)

    def set_weights(self, weights):
        """Set the weights of the optimizer.

    The weights of an optimizer are its state (ie, variables).
    This function takes the weight values associated with this
    optimizer as a list of Numpy arrays. The first value is always the
    iterations count of the optimizer, followed by the optimizer's state
    variables in the order they are created. The passed values are used to set
    the new state of the optimizer.

    For example, the RMSprop optimizer for this simple model takes a list of
    three values-- the iteration count, followed by the root-mean-square value
    of the kernel and bias of the single Dense layer:

    >>> opt = tf.keras.optimizers.RMSprop()
    >>> m = tf.keras.models.Sequential([tf.keras.layers.Dense(10)])
    >>> m.compile(opt, loss='mse')
    >>> data = np.arange(100).reshape(5, 20)
    >>> labels = np.zeros(5)
    >>> results = m.fit(data, labels)  # Training.
    >>> new_weights = [np.array(10), np.ones([20, 10]), np.zeros([10])]
    >>> opt.set_weights(new_weights)
    >>> opt.iterations
    <tf.Variable 'RMSprop/iter:0' shape=() dtype=int64, numpy=10>

    Args:
        weights: weight values as a list of numpy arrays.
    """
        params = self.weights
        if len(params) != len(weights):
            raise ValueError('You called `set_weights(weights)` on optimizer ' + self._name + ' with a  weight list of length ' + str(len(weights)) + ', but the optimizer was expecting ' + str(len(params)) + ' weights. Provided weights: ' + str(weights)[:50] + '...')
        if not params:
            return
        weight_value_tuples = []
        param_values = backend.batch_get_value(params)
        for pv, p, w in zip(param_values, params, weights):
            if pv.shape != w.shape:
                raise ValueError('Optimizer weight shape ' + str(pv.shape) + ' not compatible with provided weight shape ' + str(w.shape))
            weight_value_tuples.append((p, w))
        backend.batch_set_value(weight_value_tuples)

    def add_weight(self, name, shape, dtype=None, initializer='zeros', trainable=None, synchronization=tf_variables.VariableSynchronization.AUTO, aggregation=tf_variables.VariableAggregation.NONE):
        if dtype is None:
            dtype = dtypes.float32
        if isinstance(initializer, str) or callable(initializer):
            initializer = initializers.get(initializer)
        if synchronization == tf_variables.VariableSynchronization.ON_READ:
            if trainable:
                raise ValueError('Synchronization value can be set to VariableSynchronization.ON_READ only for non-trainable variables. You have specified trainable=True and synchronization=VariableSynchronization.ON_READ.')
            else:
                trainable = False
        elif trainable is None:
            trainable = True
        variable = self._add_variable_with_custom_getter(name=name, shape=shape, getter=base_layer_utils.make_variable, overwrite=True, initializer=initializer, dtype=dtype, trainable=trainable, use_resource=True, synchronization=synchronization, aggregation=aggregation)
        backend.track_variable(variable)
        return variable

    def _init_set_name(self, name, zero_based=True):
        if not name:
            self._name = backend.unique_object_name(generic_utils.to_snake_case(self.__class__.__name__), zero_based=zero_based)
        else:
            self._name = name

    def _assert_valid_dtypes(self, tensors):
        """Asserts tensors are all valid types (see `_valid_dtypes`).

    Args:
      tensors: Tensors to check.

    Raises:
      ValueError: If any tensor is not a valid type.
    """
        valid_dtypes = self._valid_dtypes()
        for t in tensors:
            dtype = t.dtype.base_dtype
            if dtype not in valid_dtypes:
                raise ValueError('Invalid type %r for %s, expected: %s.' % (dtype, t.name, [v for v in valid_dtypes]))

    def _valid_dtypes(self):
        """Valid types for loss, variables and gradients.

    Subclasses should override to allow other float types.

    Returns:
      Valid types for loss, variables and gradients.
    """
        return _DEFAULT_VALID_DTYPES

    def _call_if_callable(self, param):
        """Call the function if param is callable."""
        return param() if callable(param) else param

    def _resource_apply_dense(self, grad, handle, apply_state):
        """Add ops to apply dense gradients to the variable `handle`.

    Args:
      grad: a `Tensor` representing the gradient.
      handle: a `Tensor` of dtype `resource` which points to the variable to be
        updated.
      apply_state: A dict which is used across multiple apply calls.

    Returns:
      An `Operation` which updates the value of the variable.
    """
        raise NotImplementedError('Must be implemented in subclasses.')

    def _resource_apply_sparse_duplicate_indices(self, grad, handle, indices, **kwargs):
        """Add ops to apply sparse gradients to `handle`, with repeated indices.

    Optimizers which override this method must deal with repeated indices. See
    the docstring of `_apply_sparse_duplicate_indices` for details. By default
    the correct behavior, to sum non-unique indices and their associated
    gradients, is enforced by first pre-processing `grad` and `indices` and
    passing them on to `_resource_apply_sparse`. Optimizers which deal correctly
    with duplicate indices may instead override this method to avoid the
    overhead of summing.

    Args:
      grad: a `Tensor` representing the gradient for the affected indices.
      handle: a `Tensor` of dtype `resource` which points to the variable to be
        updated.
      indices: a `Tensor` of integral type representing the indices for which
        the gradient is nonzero. Indices may be repeated.
      **kwargs: May optionally contain `apply_state`

    Returns:
      An `Operation` which updates the value of the variable.
    """
        summed_grad, unique_indices = _deduplicate_indexed_slices(values=grad, indices=indices)
        return self._resource_apply_sparse(summed_grad, handle, unique_indices, **kwargs)

    def _resource_apply_sparse(self, grad, handle, indices, apply_state):
        """Add ops to apply sparse gradients to the variable `handle`.

    Similar to `_apply_sparse`, the `indices` argument to this method has been
    de-duplicated. Optimizers which deal correctly with non-unique indices may
    instead override `_resource_apply_sparse_duplicate_indices` to avoid this
    overhead.

    Args:
      grad: a `Tensor` representing the gradient for the affected indices.
      handle: a `Tensor` of dtype `resource` which points to the variable to be
        updated.
      indices: a `Tensor` of integral type representing the indices for which
        the gradient is nonzero. Indices are unique.
      apply_state: A dict which is used across multiple apply calls.

    Returns:
      An `Operation` which updates the value of the variable.
    """
        raise NotImplementedError('Must be implemented in subclasses.')

    def _resource_scatter_add(self, x, i, v):
        with ops.control_dependencies([gen_resource_variable_ops.ResourceScatterAdd(resource=x.handle, indices=i, updates=v)]):
            return x.value()

    def _resource_scatter_update(self, x, i, v):
        with ops.control_dependencies([gen_resource_variable_ops.ResourceScatterUpdate(resource=x.handle, indices=i, updates=v)]):
            return x.value()

    @property
    @layer_utils.cached_per_instance
    def _dense_apply_args(self):
        return tf_inspect.getfullargspec(self._resource_apply_dense).args

    @property
    @layer_utils.cached_per_instance
    def _sparse_apply_args(self):
        return tf_inspect.getfullargspec(self._resource_apply_sparse).args

    def _restore_slot_variable(self, slot_name, variable, slot_variable):
        """Restore a newly created slot variable's value."""
        variable_key = _var_key(variable)
        deferred_restorations = self._deferred_slot_restorations.get(slot_name, {}).pop(variable_key, [])
        deferred_restorations.sort(key=lambda position: position.restore_uid, reverse=True)
        for checkpoint_position in deferred_restorations:
            checkpoint_position.restore(slot_variable)

    def _create_or_restore_slot_variable(self, slot_variable_position, slot_name, variable):
        """Restore a slot variable's value, possibly creating it.

    Called when a variable which has an associated slot variable is created or
    restored. When executing eagerly, we create the slot variable with a
    restoring initializer.

    No new variables are created when graph building. Instead,
    _restore_slot_variable catches these after normal creation and adds restore
    ops to the graph. This method is nonetheless important when graph building
    for the case when a slot variable has already been created but `variable`
    has just been added to a dependency graph (causing us to realize that the
    slot variable needs to be restored).

    Args:
      slot_variable_position: A `trackable._CheckpointPosition` object
        indicating the slot variable `Trackable` object to be restored.
      slot_name: The name of this `Optimizer`'s slot to restore into.
      variable: The variable object this slot is being created for.
    """
        variable_key = _var_key(variable)
        slot_dict = self._slots.get(variable_key, {})
        slot_variable = slot_dict.get(slot_name, None)
        if slot_variable is None and context.executing_eagerly() and slot_variable_position.is_simple_variable() and (not ops.get_default_graph()._variable_creator_stack or self._distribution_strategy):
            initializer = trackable.CheckpointInitialValueCallable(checkpoint_position=slot_variable_position)
            slot_variable = self.add_slot(var=variable, initializer=initializer, slot_name=slot_name, shape=slot_variable_position.value_shape())
        if slot_variable is not None:
            slot_variable_position.restore(slot_variable)
        else:
            self._deferred_slot_restorations.setdefault(slot_name, {}).setdefault(variable_key, []).append(slot_variable_position)

    @contextlib.contextmanager
    def _distribution_strategy_scope(self):
        """Returns the `tf.distribute.Strategy` this optimizer was created under."""
        if self._distribution_strategy and (not distribute_lib.has_strategy()):
            with self._distribution_strategy.scope():
                yield self._distribution_strategy.scope()
        else:
            yield

Base class for Keras optimizers.

You should not use this class directly, but instead instantiate one of its
subclasses such as `tf.keras.optimizers.SGD`, `tf.keras.optimizers.Adam`, etc.

### Usage

```python
# Create an optimizer with the desired parameters.
opt = tf.keras.optimizers.SGD(learning_rate=0.1)
# `loss` is a callable that takes no argument and returns the value
# to minimize.
loss = lambda: 3 * var1 * var1 + 2 * var2 * var2
# In graph mode, returns op that minimizes the loss by updating the listed
# variables.
opt_op = opt.minimize(loss, var_list=[var1, var2])
opt_op.run()
# In eager mode, simply call minimize to update the list of variables.
opt.minimize(loss, var_list=[var1, var2])
```

### Usage in custom training loops

In Keras models, sometimes variables are created when the model is first
called, instead of construction time. Examples include 1) sequential models
without input shape pre-defined, or 2) subclassed models. Pass var_list as
callable in these cases.

Example:

```python
opt = tf.keras.optimizers.SGD(learning_rate=0.1)
model = tf.keras.Sequential()
model.add(tf.keras.layers.Dense(num_hidden, activation='relu'))
model.add(tf.keras.layers.Dense(num_classes, activation='sigmoid'))
loss_fn = lambda: tf.keras.losses.mse(model(input), output)
var_list_fn = lambda: model.trainable_weights
for input, output in data:
  opt.minimize(loss_fn, var_list_fn)
```

### Processing gradients before applying them

Calling `minimize()` takes care of both computing the gradients and
applying them to the variables.  If you want to process the gradients
before applying them you can instead use the optimizer in three steps:

1.  Compute the gradients with `tf.GradientTape`.
2.  Process the gradients as you wish.
3.  Apply the processed gradients with `apply_gradients()`.

Example:

```python
# Create an optimizer.
opt = tf.keras.optimizers.SGD(learning_rate=0.1)

# Compute the gradients for a list of variables.
with tf.GradientTape() as tape:
  loss = <call_loss_function>
vars = <list_of_variables>
grads = tape.gradient(loss, vars)

# Process the gradients, for example cap them, etc.
# capped_grads = [MyCapper(g) for g in grads]
processed_grads = [process_gradient(g) for g in grads]

# Ask the optimizer to apply the processed gradients.
opt.apply_gradients(zip(processed_grads, var_list))
```

### Use with `tf.distribute.Strategy`

This optimizer class is `tf.distribute.Strategy` aware, which means it
automatically sums gradients across all replicas. To average gradients,
you divide your loss by the global batch size, which is done
automatically if you use `tf.keras` built-in training or evaluation loops.
See the `reduction` argument of your loss which should be set to
`tf.keras.losses.Reduction.SUM_OVER_BATCH_SIZE` for averaging or
`tf.keras.losses.Reduction.SUM` for not.

To aggregate gradients yourself, call `apply_gradients` with
`experimental_aggregate_gradients` set to False. This is useful if you need to
process aggregated gradients.

If you are not using these and you want to average gradients, you should use
`tf.math.reduce_sum` to add up your per-example losses and then divide by the
global batch size. Note that when using `tf.distribute.Strategy`, the first
component of a tensor's shape is the *replica-local* batch size, which is off
by a factor equal to the number of replicas being used to compute a single
step. As a result, using `tf.math.reduce_mean` will give the wrong answer,
resulting in gradients that can be many times too big.

### Variable Constraints

All Keras optimizers respect variable constraints. If constraint function is
passed to any variable, the constraint will be applied to the variable after
the gradient has been applied to the variable.
Important: If gradient is sparse tensor, variable constraint is not supported.

### Thread Compatibility

The entire optimizer is currently thread compatible, not thread-safe. The user
needs to perform synchronization if necessary.

### Slots

Many optimizer subclasses, such as `Adam` and `Adagrad` allocate and manage
additional variables associated with the variables to train.  These are called
<i>Slots</i>.  Slots have names and you can ask the optimizer for the names of
the slots that it uses.  Once you have a slot name you can ask the optimizer
for the variable it created to hold the slot value.

This can be useful if you want to log debug a training algorithm, report stats
about the slots, etc.

### Hyperparameters

These are arguments passed to the optimizer subclass constructor
(the `__init__` method), and then passed to `self._set_hyper()`.
They can be either regular Python values (like 1.0), tensors, or
callables. If they are callable, the callable will be called during
`apply_gradients()` to get the value for the hyper parameter.

Hyperparameters can be overwritten through user code:

Example:

```python
# Create an optimizer with the desired parameters.
opt = tf.keras.optimizers.SGD(learning_rate=0.1)
# `loss` is a callable that takes no argument and returns the value
# to minimize.
loss = lambda: 3 * var1 + 2 * var2
# In eager mode, simply call minimize to update the list of variables.
opt.minimize(loss, var_list=[var1, var2])
# update learning rate
opt.learning_rate = 0.05
opt.minimize(loss, var_list=[var1, var2])
```

### Callable learning rate

Optimizer accepts a callable learning rate in two ways. The first way is
through built-in or customized
`tf.keras.optimizers.schedules.LearningRateSchedule`. The schedule will be
called on each iteration with `schedule(iteration)`, a `tf.Variable`
owned by the optimizer.

Example:

>>> var = tf.Variable(np.random.random(size=(1,)))
>>> learning_rate = tf.keras.optimizers.schedules.ExponentialDecay(
... initial_learning_rate=.01, decay_steps=20, decay_rate=.1)
>>> opt = tf.keras.optimizers.SGD(learning_rate=learning_rate)
>>> loss = lambda: 3 * var
>>> opt.minimize(loss, var_list=[var])
<tf.Variable...

The second way is through a callable function that
does not accept any arguments.

Example:

>>> var = tf.Variable(np.random.random(size=(1,)))
>>> def lr_callable():
...   return .1
>>> opt = tf.keras.optimizers.SGD(learning_rate=lr_callable)
>>> loss = lambda: 3 * var
>>> opt.minimize(loss, var_list=[var])
<tf.Variable...

### Creating a custom optimizer

If you intend to create your own optimization algorithm, simply inherit from
this class and override the following methods:

  - `_resource_apply_dense` (update variable given gradient tensor is a dense
    `tf.Tensor`)
  - `_resource_apply_sparse` (update variable given gradient tensor is a
    sparse `tf.IndexedSlices`. The most common way for this to happen
    is if you are taking the gradient through a `tf.gather`.)
  - `_create_slots`
    (if your optimizer algorithm requires additional variables)
  - `get_config`
    (serialization of the optimizer, include all hyper parameters)


def sigmoid_cross_entropy_loss(inputs: torch.Tensor, labels: torch.Tensor, num_masks: int) -> torch.Tensor:
    criterion = nn.BCEWithLogitsLoss(reduction='none')
    cross_entropy_loss = criterion(inputs, labels)
    loss = cross_entropy_loss.mean(1).sum() / num_masks
    return loss

Args:
    inputs (`torch.Tensor`):
        A float tensor of arbitrary shape.
    labels (`torch.Tensor`):
        A tensor with the same shape as inputs. Stores the binary classification labels for each element in inputs
        (0 for the negative class and 1 for the positive class).

Returns:
    loss (`torch.Tensor`): The computed loss.


def _clone_helper(op_to_clone, variant_tensor_ops):
    remap_dict = {}
    for input_tensor in op_to_clone.inputs:
        input_tensor_op = input_tensor.op
        if input_tensor_op in variant_tensor_ops:
            recursive_map = _clone_helper(input_tensor_op, variant_tensor_ops)
            remap_dict.update(recursive_map)
    inputs_list = []
    for input_tensor in op_to_clone.inputs:
        input_tensor_op = input_tensor.op
        if input_tensor_op in remap_dict:
            remapped_input = remap_dict[input_tensor_op].outputs[0]
            inputs_list.append(remapped_input)
        else:
            inputs_list.append(input_tensor_op.outputs[input_tensor.value_index])
    g = ops.get_default_graph()
    new_op = g.create_op(op_to_clone.type, inputs_list, [o.dtype for o in op_to_clone.outputs], name=op_to_clone.name, attrs=op_to_clone.node_def.attr, op_def=_get_op_def(op_to_clone))
    remap_dict[op_to_clone] = new_op
    return remap_dict

Helper method that recursively clones `op_to_clone`.

Args:
  op_to_clone: The op we want to clone.
  variant_tensor_ops: A list of ops that we have to clone along the way.

Returns:
  A dictionary mapping old_ops to new_ops created. Includes op_to_clone
  as a key.


def parse_single_example(serialized, features, name=None, example_names=None):
    return parse_single_example_v2(serialized, features, example_names, name)

Parses a single `Example` proto.

Similar to `parse_example`, except:

For dense tensors, the returned `Tensor` is identical to the output of
`parse_example`, except there is no batch dimension, the output shape is the
same as the shape given in `dense_shape`.

For `SparseTensor`s, the first (batch) column of the indices matrix is removed
(the indices matrix is a column vector), the values vector is unchanged, and
the first (`batch_size`) entry of the shape vector is removed (it is now a
single element vector).

One might see performance advantages by batching `Example` protos with
`parse_example` instead of using this function directly.

Args:
  serialized: A scalar string Tensor, a single serialized Example.
  features: A mapping of feature keys to `FixedLenFeature` or
    `VarLenFeature` values.
  name: A name for this operation (optional).
  example_names: (Optional) A scalar string Tensor, the associated name.

Returns:
  A `dict` mapping feature keys to `Tensor` and `SparseTensor` values.

Raises:
  ValueError: if any feature is invalid.


def topological_sort(nodes: Iterable[_SuccessorNode]) -> Generator[_SuccessorNode, None, None]:
    incoming = {node: set(getattr(node, 'incoming', ())) for node in nodes}
    outgoing = collections.defaultdict(set)
    for node in nodes:
        for inc in incoming[node]:
            outgoing[inc].add(node)
    stack = [node for node in nodes if not incoming[node]]
    for _ in nodes:
        if not stack:
            raise ValueError('Circular graph')
        leaf = stack.pop()
        yield leaf
        for out in outgoing[leaf]:
            incoming[out].remove(leaf)
            if not incoming[out]:
                stack.append(out)
    assert not stack

Sort a list of nodes topologically.

This will order the nodes so that any node that appears in the "incoming"
list of another node n2 will appear in the output before n2. It assumes that
the graph doesn't have any cycles.
If there are multiple ways to sort the list, a random one is picked.

Args:
  nodes: A sequence of nodes. Each node may have an attribute "incoming", a
    list of nodes (every node in this list needs to be in "nodes"). If
    "incoming" is not there, it's assumed to be empty. The list of nodes can't
    have duplicates.

Yields:
  The nodes in their topological order.
Raises:
  ValueError: If the graph contains a cycle.