def google_api_initilaize(config, api_call, alias=None):
    if api_call['function'].endswith('list') or alias == 'list':
        api_call['iterate'] = True
    if api_call['api'] == 'dfareporting':
        if not api_call['function'].startswith('userProfiles'):
            is_superuser, profile_id = get_profile_for_api(config, api_call['auth'], api_call['kwargs']['id'] if api_call['function'] == 'accounts.get' else api_call['kwargs']['accountId'])
            api_call['kwargs']['profileId'] = profile_id
            if is_superuser:
                api_call['version'] = 'prerelease'
            elif 'accountId' in api_call['kwargs']:
                del api_call['kwargs']['accountId']

Some Google API calls require a lookup or pre-call, add it here.

Modifies the API call before actual execution with any data
specifically required by an endpoint.  Currently:

  > dfa-reporting - look up user profile and add to call.

Args:
  api_call (dict): the JSON for the API call as defined in recipe.
  alias (string): mostly used to signal a list behavior (change to iterate in future?)

Returns (dict):
  A modified JSON with additional API values added.
  Currently mostly used by dfareporting API to add profile and account.

Raises:
  ValueError: If a required key in the recipe is missing.


def m_to_mile(value):
    if value is None:
        return None
    return value / 1609

Converts distance in meters to miles

Args:
  value: floating point representing the distance in meters
Returns: distance in miles


class GraniteMoeMoE(nn.Module):

    def __init__(self, config: GraniteMoeConfig):
        super(GraniteMoeMoE, self).__init__()
        self.input_size = config.hidden_size
        self.hidden_size = config.intermediate_size
        self.activation = ACT2FN[config.hidden_act]
        self.input_linear = GraniteMoeParallelExperts(config.num_local_experts, self.input_size, self.hidden_size * 2)
        self.output_linear = GraniteMoeParallelExperts(config.num_local_experts, self.hidden_size, self.input_size)
        self.router = GraniteMoeTopKGating(input_size=self.input_size, num_experts=config.num_local_experts, top_k=config.num_experts_per_tok)

    def forward(self, layer_input):
        """
        Forward pass of the mixture of experts layer.

        Args:
            layer_input (Tensor):
                Input tensor.

        Returns:
            Tensor:
                Output tensor.
            Tensor:
                Router logits.
        """
        bsz, length, emb_size = layer_input.size()
        layer_input = layer_input.reshape(-1, emb_size)
        _, batch_index, batch_gates, expert_size, router_logits = self.router(layer_input)
        expert_inputs = layer_input[batch_index]
        hidden_states = self.input_linear(expert_inputs, expert_size)
        chunked_hidden_states = hidden_states.chunk(2, dim=-1)
        hidden_states = self.activation(chunked_hidden_states[0]) * chunked_hidden_states[1]
        expert_outputs = self.output_linear(hidden_states, expert_size)
        expert_outputs = expert_outputs * batch_gates[:, None]
        zeros = torch.zeros((bsz * length, self.input_size), dtype=expert_outputs.dtype, device=expert_outputs.device)
        layer_output = zeros.index_add(0, batch_index, expert_outputs)
        layer_output = layer_output.view(bsz, length, self.input_size)
        return (layer_output, router_logits)

A Sparsely gated mixture of experts layer with 1-layer Feed-Forward networks as experts.

Args:
    config:
        Configuration object with model hyperparameters.


def connect_to_cluster(cluster_spec_or_resolver, job_name='localhost', task_index=0, protocol=None, make_master_device_default=True, cluster_device_filters=None):
    if not context.executing_eagerly():
        raise ValueError('`tf.config.experimental_connect_to_cluster` can only be called in eager mode.')
    protocol = protocol or remote_utils.get_default_communication_protocol()
    if isinstance(cluster_spec_or_resolver, server_lib.ClusterSpec):
        cluster_spec = cluster_spec_or_resolver
    elif isinstance(cluster_spec_or_resolver, cluster_resolver.ClusterResolver):
        if cluster_spec_or_resolver.master() in _LOCAL_MASTERS:
            return
        cluster_spec = cluster_spec_or_resolver.cluster_spec()
    else:
        raise ValueError('`cluster_spec_or_resolver` must be a `ClusterSpec` or a `ClusterResolver`.')
    cluster_def = copy.deepcopy(cluster_spec.as_cluster_def())
    if cluster_device_filters:
        if isinstance(cluster_device_filters, server_lib.ClusterDeviceFilters):
            cluster_device_filters = copy.deepcopy(cluster_device_filters._as_cluster_device_filters())
        else:
            raise ValueError('`cluster_device_filters` must be an instance of `tf.train.experimental.ClusterDeviceFilters`.')
    is_server_def_changed = False
    current_server_def = context.get_server_def()
    if current_server_def and job_name not in cluster_spec.jobs:
        for i, job in enumerate(current_server_def.cluster.job):
            if job.name == job_name:
                del current_server_def.cluster.job[i]
    if current_server_def is None or current_server_def.cluster != cluster_def or current_server_def.job_name != job_name or (current_server_def.task_index != task_index):
        is_server_def_changed = True
    if job_name not in cluster_spec.jobs:
        local_port = pywrap_tfe.TF_PickUnusedPortOrDie()
        job_def = cluster_def.job.add()
        job_def.name = job_name
        job_def.tasks[0] = 'localhost:{}'.format(local_port)
    if context.context().coordination_service is None:
        service_type = remote_utils.coordination_service_type(protocol)
        service_leader = ''
        if isinstance(cluster_spec_or_resolver, cluster_resolver.ClusterResolver) and hasattr(cluster_spec_or_resolver, 'tpu_hardware_feature'):
            service_leader = cluster_spec_or_resolver.get_coordination_service_leader()
            if cluster_spec_or_resolver.environment == 'google':
                is_uptc_sess = '.uptc-worker.' in cluster_spec_or_resolver.master()
                service_type = remote_utils.coordination_service_type(protocol, is_uptc_sess)
            else:
                service_type = 'standalone'
        if service_type:
            context.context().configure_coordination_service(service_type=service_type, service_leader=service_leader, enable_health_check=False)
    default_session_config = copy.deepcopy(context.context().config)
    for name in cluster_spec.jobs:
        if name == job_name:
            continue
        default_session_config.experimental.collective_group_leader = f'/job:{name}/replica:0/task:0'
    logging.info('default session config: %s', default_session_config)
    server_def = ServerDef(cluster=cluster_def, job_name=job_name, task_index=task_index, protocol=protocol, default_session_config=default_session_config, cluster_device_filters=cluster_device_filters)
    if is_server_def_changed:
        context.set_server_def(server_def)
    else:
        context.update_server_def(server_def)
    if make_master_device_default and isinstance(cluster_spec_or_resolver, cluster_resolver.ClusterResolver) and cluster_spec_or_resolver.master():
        master = cluster_spec_or_resolver.master()
        master_job_name = None
        master_task_id = None
        for job_name in cluster_spec.jobs:
            for task_id in cluster_spec.task_indices(job_name):
                task_address = cluster_spec.task_address(job_name, task_id)
                if master in task_address or task_address in master:
                    master_job_name = job_name
                    master_task_id = task_id
                    break
        if not master_job_name:
            raise ValueError('`make_master_device_default` is set to True but cannot find master %s in the cluster' % master)
        master_device = '/job:{}/replica:0/task:{}'.format(master_job_name, master_task_id)
        master_device = device_util.canonicalize(master_device)
        current_device = device_util.current()
        if current_device:
            current_device = device_util.canonicalize(current_device)
        if current_device and current_device != master_device:
            raise ValueError('`connect_to_cluster` is called inside existing device scope %s, which is different from the master device scope %s to enter. This is not allowed.' % (current_device, master_device))
        if not current_device:
            logging.info('Entering into master device scope: %s', master_device)
            ops.device(master_device).__enter__()

Connects to the given cluster.

Will make devices on the cluster available to use. Note that calling this more
than once will work, but will invalidate any tensor handles on the old remote
devices.

If the given local job name is not present in the cluster specification, it
will be automatically added, using an unused port on the localhost.

Device filters can be specified to isolate groups of remote tasks to avoid
undesired accesses between workers. Workers accessing resources or launching
ops / functions on filtered remote devices will result in errors (unknown
devices). For any remote task, if no device filter is present, all cluster
devices will be visible; if any device filter is specified, it can only
see devices matching at least one filter. Devices on the task itself are
always visible. Device filters can be particially specified.

For example, for a cluster set up for parameter server training, the following
device filters might be specified:

```python
cdf = tf.config.experimental.ClusterDeviceFilters()
# For any worker, only the devices on PS nodes and itself are visible
for i in range(num_workers):
  cdf.set_device_filters('worker', i, ['/job:ps'])
# Similarly for any ps, only the devices on workers and itself are visible
for i in range(num_ps):
  cdf.set_device_filters('ps', i, ['/job:worker'])

tf.config.experimental_connect_to_cluster(cluster_def,
                                          cluster_device_filters=cdf)
```

Args:
  cluster_spec_or_resolver: A `ClusterSpec` or `ClusterResolver` describing
    the cluster.
  job_name: The name of the local job.
  task_index: The local task index.
  protocol: The communication protocol, such as `"grpc"`. If unspecified, will
    use the default from `python/platform/remote_utils.py`.
  make_master_device_default: If True and a cluster resolver is passed, will
    automatically enter the master task device scope, which indicates the
    master becomes the default device to run ops. It won't do anything if
    a cluster spec is passed. Will throw an error if the caller is currently
    already in some device scope.
  cluster_device_filters: an instance of
    `tf.train.experimental/ClusterDeviceFilters` that specify device filters
    to the remote tasks in cluster.


def convert_graphdef_with_arrays(input_data, input_arrays_with_shape, output_arrays, control_output_arrays, **kwargs):
    model_flags = build_model_flags(**kwargs)
    conversion_flags = build_conversion_flags(**kwargs)
    quantized_input_stats = kwargs.get('quantized_input_stats', None)
    for idx, (name, shape) in enumerate(input_arrays_with_shape):
        input_array = model_flags.input_arrays.add()
        if _is_quantized_input_stats_required(conversion_flags):
            if quantized_input_stats:
                input_array.mean_value, input_array.std_value = quantized_input_stats[idx]
            else:
                raise ValueError('The `quantized_input_stats` flag must be defined when either `inference_type` flag or `inference_input_type` flag is set to tf.int8 or tf.uint8.')
        input_array.name = name
        input_array.shape.dims.extend(list(map(int, shape)))
    if output_arrays:
        for name in output_arrays:
            model_flags.output_arrays.append(name)
    if control_output_arrays:
        for name in control_output_arrays:
            model_flags.control_output_arrays.append(name)
    data = convert(model_flags, conversion_flags, input_data.SerializeToString(), debug_info_str=None)
    return data

Convert a frozen GraphDef that can't be loaded in TF.

Conversion can be customized by providing arguments that are forwarded to
`build_model_flags` and `build_conversion_flags` (see documentation).

Args:
  input_data: Input data (i.e. often `sess.graph_def`),
  input_arrays_with_shape: Tuple of strings representing input tensor names
    and list of integers representing input shapes (e.g., [("foo" : [1, 16,
    16, 3])]). Use only when graph cannot be loaded into TensorFlow and when
    `input_tensors` is None.
  output_arrays: List of output tensors to freeze graph with. Use only when
    graph cannot be loaded into TensorFlow and when `output_tensors` is None.
  control_output_arrays: Control output node names. This is used when
    converting a Graph with no output tensors. For example, if the graph's
    last operation is a Print op, just specify that op's name in this field.
    This can be used together with the `output_arrays` parameter.
  **kwargs: See `build_model_flags` and `build_conversion_flags`.

Returns:
  The converted data. For example if TFLite was the destination, then
  this will be a tflite flatbuffer in a bytes array.

Raises:
  Defined in `build_conversion_flags`.


def evaluate(code: str, *, global_vars: Optional[Dict[str, Any]]=None, permission: Optional[permissions.CodePermission]=None, returns_stdout: bool=False, outputs_intermediate: bool=False) -> Union[Any, Dict[str, Any]]:
    permission = permission or permissions.get_permission()
    ctx = dict(get_context())
    if global_vars:
        ctx.update(global_vars)
    code_block = parsing.parse(code, permission)
    global_vars, orig_global_vars = (ctx, ctx.copy())
    if not code_block.body:
        return {} if outputs_intermediate else None
    stdout = io.StringIO()
    with contextlib.redirect_stdout(stdout):
        if hasattr(code_block.body[-1], 'value'):
            last_expr = code_block.body.pop()
            result_vars = [RESULT_KEY]
            if isinstance(last_expr, ast.Assign):
                for name_node in last_expr.targets:
                    if isinstance(name_node, ast.Name):
                        result_vars.append(name_node.id)
            last_expr = ast.Expression(last_expr.value)
            try:
                exec(compile(code_block, '', mode='exec'), global_vars)
                result = eval(compile(last_expr, '', mode='eval'), global_vars)
            except BaseException as e:
                raise errors.CodeError(code, e) from e
            for result_var in result_vars:
                global_vars[result_var] = result
        else:
            try:
                exec(compile(code_block, '', mode='exec'), global_vars)
            except BaseException as e:
                raise errors.CodeError(code, e) from e
            global_vars[RESULT_KEY] = list(global_vars.values())[-1]
    if returns_stdout:
        return stdout.getvalue()
    if outputs_intermediate:
        outputs = {}
        for k, v in global_vars.items():
            if k == '__builtins__':
                continue
            if k not in orig_global_vars or v is not orig_global_vars[k]:
                outputs[k] = v
        outputs[STDOUT_KEY] = stdout.getvalue()
        return outputs
    return global_vars[RESULT_KEY]

Executes Python code.

Features:
  * Fine-grained execution policy for limiting what APIs could be executed.
    This eliminates the need for sandboxing.
  * It exposes both the final results and intermediate results (variables).

Args:
  code: Python code to run.
  global_vars: An optional dict as the globals that could be referenced by the
    code.
  permission: Permission for the Python code to run.
  returns_stdout: If True, the stdout (a str) will be returned.
  outputs_intermediate: Applicable when returns_stdout is False. If True,
    intermediate output will be outputted as a dict, with the last line's
    value accessible by key '__result__' and the std output accessible by
    key '__stdout__'. Otherwise the value of the last line will be returned.

Returns:
  The value of the last line of the code block. Or a dict of variable
  names of all locals to their evaluated values as the output of the code to
  run. The value for the last line can be accessed by key '__result__'. Or the
  stdout as a str.


def run_change_point_analysis(test_config_container: TestConfigContainer, big_query_metrics_fetcher: MetricsFetcher, change_point_config: ChangePointConfig=ChangePointConfig(), save_alert_metadata: bool=False):
    logging.info('Running change point analysis for test ID :%s on metric: % s' % (test_config_container.test_id, test_config_container.metric_name))
    test_name = test_config_container.test_name
    min_runs_between_change_points = change_point_config.min_runs_between_change_points
    num_runs_in_change_point_window = change_point_config.num_runs_in_change_point_window
    metric_container = big_query_metrics_fetcher.fetch_metric_data(test_config=test_config_container)
    metric_container.sort_by_timestamp()
    metric_values = metric_container.values
    timestamps = metric_container.timestamps
    change_point_index = find_latest_change_point_index(metric_values=metric_values)
    if not change_point_index:
        logging.info('Change point is not detected for the test ID %s' % test_config_container.test_id)
        return False
    latest_change_point_run = len(timestamps) - 1 - change_point_index
    if not is_change_point_in_valid_window(num_runs_in_change_point_window, latest_change_point_run):
        logging.info('Performance regression/improvement found for the test ID: %s. on metric %s. Since the change point run %s lies outside the num_runs_in_change_point_window distance: %s, alert is not raised.' % (test_config_container.test_id, test_config_container.metric_name, latest_change_point_run + 1, num_runs_in_change_point_window))
        return False
    is_valid_change_point = True
    last_reported_issue_number = None
    issue_metadata_table_name = f'{test_config_container.metrics_table}_{test_config_container.metric_name}'
    if test_config_container.test_name:
        issue_metadata_table_name = f'{issue_metadata_table_name}_{test_config_container.test_name}'
    existing_issue_data = get_existing_issues_data(table_name=issue_metadata_table_name)
    if existing_issue_data is not None:
        existing_issue_timestamps = existing_issue_data[constants._CHANGE_POINT_TIMESTAMP_LABEL].tolist()
        last_reported_issue_number = existing_issue_data[constants._ISSUE_NUMBER].tolist()[0]
        if not isinstance(last_reported_issue_number, int):
            last_reported_issue_number = last_reported_issue_number.item()
        is_valid_change_point = is_sibling_change_point(previous_change_point_timestamps=existing_issue_timestamps, change_point_index=change_point_index, timestamps=timestamps, min_runs_between_change_points=min_runs_between_change_points, test_id=test_config_container.test_id)
    if is_valid_change_point and save_alert_metadata:
        issue_number, issue_url = create_performance_alert(test_config_container=test_config_container, metric_container=metric_container, change_point_index=change_point_index, existing_issue_number=last_reported_issue_number)
        issue_metadata = GitHubIssueMetaData(issue_timestamp=pd.Timestamp(datetime.now().replace(tzinfo=timezone.utc)), test_id=test_config_container.test_id.replace('.', '_'), test_name=test_name or uuid.uuid4().hex, metric_name=test_config_container.metric_name, change_point=metric_values[change_point_index], issue_number=issue_number, issue_url=issue_url, change_point_timestamp=timestamps[change_point_index])
        publish_issue_metadata_to_big_query(issue_metadata=issue_metadata, table_name=issue_metadata_table_name, project=test_config_container.project)
    return is_valid_change_point

Args:
 test_config_container: TestConfigContainer containing test metadata for
  fetching data and running change point analysis.
 big_query_metrics_fetcher: BigQuery metrics fetcher used to fetch data for
  change point analysis.
  change_point_config: ChangePointConfig containing parameters to run
    change point analysis.
  save_alert_metadata: bool indicating if issue metadata
    should be published to BigQuery table.
Returns:
   bool indicating if a change point is observed and alerted on GitHub.


def get_doctest_files(diff_with_last_commit: bool=False) -> List[str]:
    repo = Repo(PATH_TO_REPO)
    test_files_to_run = []
    if not diff_with_last_commit:
        print(f'main is at {repo.refs.main.commit}')
        print(f'Current head is at {repo.head.commit}')
        branching_commits = repo.merge_base(repo.refs.main, repo.head)
        for commit in branching_commits:
            print(f'Branching commit: {commit}')
        test_files_to_run = get_diff_for_doctesting(repo, repo.head.commit, branching_commits)
    else:
        print(f'main is at {repo.head.commit}')
        parent_commits = repo.head.commit.parents
        for commit in parent_commits:
            print(f'Parent commit: {commit}')
        test_files_to_run = get_diff_for_doctesting(repo, repo.head.commit, parent_commits)
    all_test_files_to_run = get_all_doctest_files()
    new_test_files = get_new_doctest_files(repo, repo.head.commit, repo.refs.main.commit)
    test_files_to_run = list(set(test_files_to_run + new_test_files))
    with open('utils/slow_documentation_tests.txt') as fp:
        slow_documentation_tests = set(fp.read().strip().split('\n'))
    test_files_to_run = [x for x in test_files_to_run if x in all_test_files_to_run and x not in slow_documentation_tests]
    test_files_to_run = [f for f in test_files_to_run if (PATH_TO_REPO / f).exists()]
    return sorted(test_files_to_run)

Return a list of python and Markdown files where doc example have been modified between:

- the current head and the main branch if `diff_with_last_commit=False` (default)
- the current head and its parent commit otherwise.

Returns:
    `List[str]`: The list of Python and Markdown files with a diff (files added or renamed are always returned, files
    modified are returned if the diff in the file is only in doctest examples).


def GetRawKeyFunction():
    for get_raw_key_function in (_GetRawKeyFunctionPosix, _GetRawKeyFunctionWindows):
        try:
            return get_raw_key_function()
        except:
            pass
    return lambda: None

Returns a function that reads one keypress from stdin with no echo.

Returns:
  A function that reads one keypress from stdin with no echo or a function
  that always returns None if stdin does not support it.


def assert_type(tensor, tf_type, message=None, name=None):
    tf_type = dtypes.as_dtype(tf_type)
    with ops.name_scope(name, 'assert_type', [tensor]):
        if not isinstance(tensor, sparse_tensor.SparseTensor):
            tensor = ops.convert_to_tensor(tensor, name='tensor')
        if tensor.dtype != tf_type:
            raise TypeError(f'{_message_prefix(message)}{getattr(tensor, 'name', 'tensor')} must be of type {tf_type!r}; got {tensor.dtype!r}')
        return control_flow_ops.no_op('statically_determined_correct_type')

Statically asserts that the given `Tensor` is of the specified type.

Args:
  tensor: A `Tensor` or `SparseTensor`.
  tf_type: A tensorflow type (`dtypes.float32`, `tf.int64`, `dtypes.bool`,
    etc).
  message: A string to prefix to the default message.
  name:  A name to give this `Op`.  Defaults to "assert_type"

Raises:
  TypeError: If the tensors data type doesn't match `tf_type`.

Returns:
  A `no_op` that does nothing.  Type can be determined statically.


def _slot_dict(self, slot_name):
    named_slots = self._slots.get(slot_name, None)
    if named_slots is None:
        named_slots = {}
        self._slots[slot_name] = named_slots
    return named_slots

Returns a dict for caching slots created under the given name.

Args:
  slot_name: Name for the slot.

Returns:
  A dict that maps primary `Variable` objects to the slot created
  for that variable, under the given slot name.


def Get(self, request, global_params=None):
    config = self.GetMethodConfig('Get')
    return self._RunMethod(config, request, global_params=global_params)

Gets information about a snapshot.

Args:
  request: (DataflowProjectsSnapshotsGetRequest) input message
  global_params: (StandardQueryParameters, default: None) global arguments
Returns:
  (Snapshot) The response message.


def assert_raises(expected_exception, extras=None, *args, **kwargs):
    context = _AssertRaisesContext(expected_exception, extras=extras)
    return context

Assert that an exception is raised when a function is called.

If no exception is raised, test fail. If an exception is raised but not
of the expected type, the exception is let through.

This should only be used as a context manager:
  with assert_raises(Exception):
    func()

Args:
  expected_exception: An exception class that is expected to be
    raised.
  extras: An optional field for extra information to be included in
    test result.


class CategoricalFocalCrossentropy(LossFunctionWrapper):

    def __init__(self, alpha=0.25, gamma=2.0, from_logits=False, label_smoothing=0.0, axis=-1, reduction='sum_over_batch_size', name='categorical_focal_crossentropy', dtype=None):
        """Initializes `CategoricalFocalCrossentropy` instance."""
        super().__init__(categorical_focal_crossentropy, name=name, reduction=reduction, dtype=dtype, alpha=alpha, gamma=gamma, from_logits=from_logits, label_smoothing=label_smoothing, axis=axis)
        self.from_logits = from_logits
        self.label_smoothing = label_smoothing
        self.axis = axis
        self.alpha = alpha
        self.gamma = gamma

    def get_config(self):
        config = Loss.get_config(self)
        config.update({'from_logits': self.from_logits, 'label_smoothing': self.label_smoothing, 'axis': self.axis, 'alpha': self.alpha, 'gamma': self.gamma})
        return config

Computes the alpha balanced focal crossentropy loss.

Use this crossentropy loss function when there are two or more label
classes and if you want to handle class imbalance without using
`class_weights`. We expect labels to be provided in a `one_hot`
representation.

According to [Lin et al., 2018](https://arxiv.org/pdf/1708.02002.pdf), it
helps to apply a focal factor to down-weight easy examples and focus more on
hard examples. The general formula for the focal loss (FL)
is as follows:

`FL(p_t) = (1 - p_t) ** gamma * log(p_t)`

where `p_t` is defined as follows:
`p_t = output if y_true == 1, else 1 - output`

`(1 - p_t) ** gamma` is the `modulating_factor`, where `gamma` is a focusing
parameter. When `gamma` = 0, there is no focal effect on the cross entropy.
`gamma` reduces the importance given to simple examples in a smooth manner.

The authors use alpha-balanced variant of focal loss (FL) in the paper:
`FL(p_t) = -alpha * (1 - p_t) ** gamma * log(p_t)`

where `alpha` is the weight factor for the classes. If `alpha` = 1, the
loss won't be able to handle class imbalance properly as all
classes will have the same weight. This can be a constant or a list of
constants. If alpha is a list, it must have the same length as the number
of classes.

The formula above can be generalized to:
`FL(p_t) = alpha * (1 - p_t) ** gamma * CrossEntropy(y_true, y_pred)`

where minus comes from `CrossEntropy(y_true, y_pred)` (CE).

Extending this to multi-class case is straightforward:
`FL(p_t) = alpha * (1 - p_t) ** gamma * CategoricalCE(y_true, y_pred)`

In the snippet below, there is `num_classes` floating pointing values per
example. The shape of both `y_pred` and `y_true` are
`(batch_size, num_classes)`.

Args:
    alpha: A weight balancing factor for all classes, default is `0.25` as
        mentioned in the reference. It can be a list of floats or a scalar.
        In the multi-class case, alpha may be set by inverse class
        frequency by using `compute_class_weight` from `sklearn.utils`.
    gamma: A focusing parameter, default is `2.0` as mentioned in the
        reference. It helps to gradually reduce the importance given to
        simple (easy) examples in a smooth manner.
    from_logits: Whether `output` is expected to be a logits tensor. By
        default, we consider that `output` encodes a probability
        distribution.
    label_smoothing: Float in [0, 1]. When > 0, label values are smoothed,
        meaning the confidence on label values are relaxed. For example, if
        `0.1`, use `0.1 / num_classes` for non-target labels and
        `0.9 + 0.1 / num_classes` for target labels.
    axis: The axis along which to compute crossentropy (the features
        axis). Defaults to `-1`.
    reduction: Type of reduction to apply to the loss. In almost all cases
        this should be `"sum_over_batch_size"`. Supported options are
        `"sum"`, `"sum_over_batch_size"`, `"mean"`,
        `"mean_with_sample_weight"` or `None`. `"sum"` sums the loss,
        `"sum_over_batch_size"` and `"mean"` sum the loss and divide by the
        sample size, and `"mean_with_sample_weight"` sums the loss and
        divides by the sum of the sample weights. `"none"` and `None`
        perform no aggregation. Defaults to `"sum_over_batch_size"`.
    name: Optional name for the loss instance.
    dtype: The dtype of the loss's computations. Defaults to `None`, which
        means using `keras.backend.floatx()`. `keras.backend.floatx()` is a
        `"float32"` unless set to different value
        (via `keras.backend.set_floatx()`). If a `keras.DTypePolicy` is
        provided, then the `compute_dtype` will be utilized.

Examples:

Standalone usage:

>>> y_true = [[0., 1., 0.], [0., 0., 1.]]
>>> y_pred = [[0.05, 0.95, 0], [0.1, 0.8, 0.1]]
>>> # Using 'auto'/'sum_over_batch_size' reduction type.
>>> cce = keras.losses.CategoricalFocalCrossentropy()
>>> cce(y_true, y_pred)
0.23315276

>>> # Calling with 'sample_weight'.
>>> cce(y_true, y_pred, sample_weight=np.array([0.3, 0.7]))
0.1632

>>> # Using 'sum' reduction type.
>>> cce = keras.losses.CategoricalFocalCrossentropy(
...     reduction="sum")
>>> cce(y_true, y_pred)
0.46631

>>> # Using 'none' reduction type.
>>> cce = keras.losses.CategoricalFocalCrossentropy(
...     reduction=None)
>>> cce(y_true, y_pred)
array([3.2058331e-05, 4.6627346e-01], dtype=float32)

Usage with the `compile()` API:

```python
model.compile(optimizer='adam',
              loss=keras.losses.CategoricalFocalCrossentropy())
```


def _prep_debug_cli_for_run_end(self, debug_dump, tf_error, passed_filter, passed_filter_exclude_node_names):
    if tf_error:
        help_intro = cli_shared.get_error_intro(tf_error)
        self._init_command = 'help'
        self._title_color = 'red_on_white'
    else:
        help_intro = None
        self._init_command = 'lt'
        self._title_color = 'black_on_white'
        if passed_filter is not None:
            self._init_command = 'lt -f %s' % passed_filter
            if passed_filter_exclude_node_names:
                self._init_command += ' --filter_exclude_node_names %s' % passed_filter_exclude_node_names
            self._title_color = 'red_on_white'
    self._run_cli = analyzer_cli.create_analyzer_ui(debug_dump, self._tensor_filters, ui_type=self._ui_type, on_ui_exit=self._remove_dump_root, config=self._config)
    dumped_tensor_names = []
    for datum in debug_dump.dumped_tensor_data:
        dumped_tensor_names.append('%s:%d' % (datum.node_name, datum.output_slot))
    self._run_cli.register_tab_comp_context(['print_tensor', 'pt'], dumped_tensor_names)
    self._run_cli.register_tab_comp_context(['node_info', 'ni', 'list_inputs', 'li', 'list_outputs', 'lo'], [str(node_name) for node_name in debug_dump.nodes()])
    self._title = 'run-end: ' + self._run_description
    if help_intro:
        self._run_cli.set_help_intro(help_intro)

Prepare (but not launch) CLI for run-end, with debug dump from the run.

Args:
  debug_dump: (debug_data.DebugDumpDir) The debug dump directory from this
    run.
  tf_error: (None or OpError) OpError that happened during the run() call
    (if any).
  passed_filter: (None or str) Name of the tensor filter that just passed
    and caused the preparation of this run-end CLI (if any).
  passed_filter_exclude_node_names: (None or str) Regular expression used
    with the tensor filter to exclude ops with names matching the regular
    expression.


def set_random_seed(seed):
    if context.executing_eagerly():
        context.set_global_seed(seed)
    else:
        ops.get_default_graph().seed = seed

Sets the graph-level random seed for the default graph.

Operations that rely on a random seed actually derive it from two seeds:
the graph-level and operation-level seeds. This sets the graph-level seed.

Its interactions with operation-level seeds is as follows:

  1. If neither the graph-level nor the operation seed is set:
    A random seed is used for this op.
  2. If the graph-level seed is set, but the operation seed is not:
    The system deterministically picks an operation seed in conjunction with
    the graph-level seed so that it gets a unique random sequence. Within the
    same version of tensorflow and user code, this sequence is deterministic.
    However across different versions, this sequence might change. If the
    code depends on particular seeds to work, specify both graph-level
    and operation-level seeds explicitly.
  3. If the graph-level seed is not set, but the operation seed is set:
    A default graph-level seed and the specified operation seed are used to
    determine the random sequence.
  4. If both the graph-level and the operation seed are set:
    Both seeds are used in conjunction to determine the random sequence.

To illustrate the user-visible effects, consider these examples:

To generate different sequences across sessions, set neither
graph-level nor op-level seeds:

```python
a = tf.random.uniform([1])
b = tf.random.normal([1])

print("Session 1")
with tf.compat.v1.Session() as sess1:
  print(sess1.run(a))  # generates 'A1'
  print(sess1.run(a))  # generates 'A2'
  print(sess1.run(b))  # generates 'B1'
  print(sess1.run(b))  # generates 'B2'

print("Session 2")
with tf.compat.v1.Session() as sess2:
  print(sess2.run(a))  # generates 'A3'
  print(sess2.run(a))  # generates 'A4'
  print(sess2.run(b))  # generates 'B3'
  print(sess2.run(b))  # generates 'B4'
```

To generate the same repeatable sequence for an op across sessions, set the
seed for the op:

```python
a = tf.random.uniform([1], seed=1)
b = tf.random.normal([1])

# Repeatedly running this block with the same graph will generate the same
# sequence of values for 'a', but different sequences of values for 'b'.
print("Session 1")
with tf.compat.v1.Session() as sess1:
  print(sess1.run(a))  # generates 'A1'
  print(sess1.run(a))  # generates 'A2'
  print(sess1.run(b))  # generates 'B1'
  print(sess1.run(b))  # generates 'B2'

print("Session 2")
with tf.compat.v1.Session() as sess2:
  print(sess2.run(a))  # generates 'A1'
  print(sess2.run(a))  # generates 'A2'
  print(sess2.run(b))  # generates 'B3'
  print(sess2.run(b))  # generates 'B4'
```

To make the random sequences generated by all ops be repeatable across
sessions, set a graph-level seed:

```python
tf.compat.v1.random.set_random_seed(1234)
a = tf.random.uniform([1])
b = tf.random.normal([1])

# Repeatedly running this block with the same graph will generate the same
# sequences of 'a' and 'b'.
print("Session 1")
with tf.compat.v1.Session() as sess1:
  print(sess1.run(a))  # generates 'A1'
  print(sess1.run(a))  # generates 'A2'
  print(sess1.run(b))  # generates 'B1'
  print(sess1.run(b))  # generates 'B2'

print("Session 2")
with tf.compat.v1.Session() as sess2:
  print(sess2.run(a))  # generates 'A1'
  print(sess2.run(a))  # generates 'A2'
  print(sess2.run(b))  # generates 'B1'
  print(sess2.run(b))  # generates 'B2'
```

@compatibility(TF2)
'tf.compat.v1.set_random_seed' is compatible with eager mode. However,
in eager mode this API will set the global seed instead of the
graph-level seed of the default graph. In TF2 this API is changed to
[tf.random.set_seed]
(https://www.tensorflow.org/api_docs/python/tf/random/set_seed).
@end_compatibility

Args:
  seed: integer.


def run_inference(self, batch, model, inference_args=None):
    _validate_inference_args(inference_args)
    vectorized_batch = np.vstack(batch)
    predictions = hdbscan.approximate_predict(model, vectorized_batch)
    return [PredictionResult(x, y) for x, y in zip(batch, predictions)]

Runs inferences on a batch of numpy arrays.

Args:
  batch: A sequence of examples as numpy arrays. They should
    be single examples.
  model: A numpy model or pipeline. Must implement predict(X).
    Where the parameter X is a numpy array.
  inference_args: Any additional arguments for an inference.

Returns:
  An Iterable of type PredictionResult.


def _check_state_for_finalize_write(self, writer_results, num_shards):
    if not writer_results:
        return ([], [], [], 0)
    src_glob = FileSystems.join(FileSystems.split(writer_results[0])[0], '*')
    dst_glob = self._get_final_name_glob(num_shards)
    src_glob_files = set((file_metadata.path for mr in FileSystems.match([src_glob]) for file_metadata in mr.metadata_list))
    dst_glob_files = set((file_metadata.path for mr in FileSystems.match([dst_glob]) for file_metadata in mr.metadata_list))
    src_files = []
    dst_files = []
    delete_files = []
    num_skipped = 0
    for shard_num, src in enumerate(writer_results):
        final_name = self._get_final_name(shard_num, num_shards)
        dst = final_name
        src_exists = src in src_glob_files
        dst_exists = dst in dst_glob_files
        if not src_exists and (not dst_exists):
            raise BeamIOError('src and dst files do not exist. src: %s, dst: %s' % (src, dst))
        if not src_exists and dst_exists:
            _LOGGER.debug('src: %s -> dst: %s already renamed, skipping', src, dst)
            num_skipped += 1
            continue
        if src_exists and dst_exists and (FileSystems.checksum(src) == FileSystems.checksum(dst)):
            _LOGGER.debug('src: %s == dst: %s, deleting src', src, dst)
            delete_files.append(src)
            continue
        src_files.append(src)
        dst_files.append(dst)
    self._report_sink_lineage(dst_glob, dst_files)
    return (src_files, dst_files, delete_files, num_skipped)

Checks writer output files' states.

Returns:
  src_files, dst_files: Lists of files to rename. For each i, finalize_write
    should rename(src_files[i], dst_files[i]).
  delete_files: Src files to delete. These could be leftovers from an
    incomplete (non-atomic) rename operation.
  num_skipped: Tally of writer results files already renamed, such as from
    a previous run of finalize_write().


class SiglipVisionConfig(PretrainedConfig):
    model_type = 'siglip_vision_model'
    base_config_key = 'vision_config'

    def __init__(self, hidden_size=768, intermediate_size=3072, num_hidden_layers=12, num_attention_heads=12, num_channels=3, image_size=224, patch_size=16, hidden_act='gelu_pytorch_tanh', layer_norm_eps=1e-06, attention_dropout=0.0, **kwargs):
        super().__init__(**kwargs)
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.num_channels = num_channels
        self.patch_size = patch_size
        self.image_size = image_size
        self.attention_dropout = attention_dropout
        self.layer_norm_eps = layer_norm_eps
        self.hidden_act = hidden_act

This is the configuration class to store the configuration of a [`SiglipVisionModel`]. It is used to instantiate a
Siglip vision encoder according to the specified arguments, defining the model architecture. Instantiating a
configuration with the defaults will yield a similar configuration to that of the vision encoder of the Siglip
[google/siglip-base-patch16-224](https://huggingface.co/google/siglip-base-patch16-224) 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.
    intermediate_size (`int`, *optional*, defaults to 3072):
        Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
    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.
    num_channels (`int`, *optional*, defaults to 3):
        Number of channels in the input images.
    image_size (`int`, *optional*, defaults to 224):
        The size (resolution) of each image.
    patch_size (`int`, *optional*, defaults to 16):
        The size (resolution) of each patch.
    hidden_act (`str` or `function`, *optional*, defaults to `"gelu_pytorch_tanh"`):
        The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
        `"relu"`, `"selu"` and `"gelu_new"` `"quick_gelu"` are supported.
    layer_norm_eps (`float`, *optional*, defaults to 1e-06):
        The epsilon used by the layer normalization layers.
    attention_dropout (`float`, *optional*, defaults to 0.0):
        The dropout ratio for the attention probabilities.

Example:

```python
>>> from transformers import SiglipVisionConfig, SiglipVisionModel

>>> # Initializing a SiglipVisionConfig with google/siglip-base-patch16-224 style configuration
>>> configuration = SiglipVisionConfig()

>>> # Initializing a SiglipVisionModel (with random weights) from the google/siglip-base-patch16-224 style configuration
>>> model = SiglipVisionModel(configuration)

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


def __init__(self, api_key: str, config: interfaces.Config | None=None):
    self._config = config or interfaces.Config()
    self._genai_processor = genai_model.GenaiModel(api_key=api_key, model_name=self._config.topic_researcher_model_name, generate_content_config=types.GenerateContentConfig(tools=self._config.enabled_research_tools))
    p_preamble = preamble.Preamble(content=[ProcessorPart(prompts.TOPIC_RESEARCH_PREAMBLE), ProcessorPart('Topic to research: ')])
    p_verbalizer = topic_verbalizer.TopicVerbalizer(config=self._config)
    p_suffix = preamble.Suffix(content=[ProcessorPart('Your research: ')])
    self._pipeline = p_verbalizer + p_preamble + p_suffix + self._genai_processor

Initializes the TopicResearcher.

Args:
  api_key: The API key to use for the GenAI API.
  config: The agent configuration.


Whether this trigger creates an empty pane even if there are no elements.

Returns:
  True if this trigger guarantees that there will always be an ON_TIME pane
  even if there are no elements in that pane.


def parse_arguments(argv):
    parser = argparse.ArgumentParser(description='online-clustering')
    parser.add_argument('-m', '--mode', help='Mode to run pipeline in.', choices=['local', 'cloud'], default='local')
    parser.add_argument('-p', '--project', help='GCP project to run pipeline on.', default=cfg.PROJECT_ID)
    args, _ = parser.parse_known_args(args=argv)
    return args

It parses the arguments passed to the command line and returns them as an object

Args:
  argv: The arguments passed to the command line.

Returns:
  The arguments that are being passed in.


class TFDebertaV2XSoftmax(keras.layers.Layer):

    def __init__(self, axis=-1, **kwargs):
        super().__init__(**kwargs)
        self.axis = axis

    def call(self, inputs: tf.Tensor, mask: tf.Tensor):
        rmask = tf.logical_not(tf.cast(mask, tf.bool))
        output = tf.where(rmask, tf.cast(float('-inf'), dtype=self.compute_dtype), inputs)
        output = stable_softmax(tf.cast(output, dtype=tf.float32), self.axis)
        output = tf.where(rmask, 0.0, output)
        return output

Masked Softmax which is optimized for saving memory

Args:
    input (`tf.Tensor`): The input tensor that will apply softmax.
    mask (`tf.Tensor`): The mask matrix where 0 indicate that element will be ignored in the softmax calculation.
    dim (int): The dimension that will apply softmax


def read(self, index, name=None):
    return self._implementation.read(index, name=name)

Read the value at location `index` in the TensorArray.

Args:
  index: 0-D.  int32 tensor with the index to read from.
  name: A name for the operation (optional).

Returns:
  The tensor at index `index`.


def download_and_prep_data() -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    mnist_dataset = tf.keras.datasets.mnist
    (tr_x, tr_y), (te_x, te_y) = mnist_dataset.load_data()
    tr_x = tr_x / 255.0
    te_x = te_x / 255.0
    return (tr_x, tr_y, te_x, te_y)

Download dataset and scale to [0, 1].

Returns:
  tr_x: Training data.
  tr_y: Training labels.
  te_x: Testing data.
  te_y: Testing labels.


def seek(self, offset, whence=os.SEEK_SET):
    self._checkClosed()
    if whence == os.SEEK_SET:
        self._position = offset
    elif whence == os.SEEK_CUR:
        self._position += offset
    elif whence == os.SEEK_END:
        self._position = self._downloader.size + offset
    else:
        raise ValueError('Whence mode %r is invalid.' % whence)
    self._position = min(self._position, self._downloader.size)
    self._position = max(self._position, 0)
    return self._position

Set the stream's current offset.

Note if the new offset is out of bound, it is adjusted to either 0 or EOF.

Args:
  offset: seek offset as number.
  whence: seek mode. Supported modes are os.SEEK_SET (absolute seek),
    os.SEEK_CUR (seek relative to the current position), and os.SEEK_END
    (seek relative to the end, offset should be negative).

Raises:
  ``ValueError``: When this stream is closed or if whence is invalid.


def recombine(self, parents: List[pg.DNA], global_state: pg.geno.AttributeDict, step: int) -> List[pg.DNA]:

Generate a list of child DNA based on the list of parents given.

User should override this method with optional keyword arguments
'global_state' and 'step'.

The parents DNA contains a metadata field 'generation', which is the
generation of the parent DNA. If the Recombinator does not assign this
field for the new child DNA, the child DNA will have the maximum generation
from the parents plus 1.

Args:
  parents: Parent trials.
  global_state: An `AttributeDict` object as the global state container,
    which is readable/writable during the operation.
  step: Number of examples historically proposed, which can be used for
    determining a cross over schedule.

Returns:
  A list of generated child DNA.


def __init__(self, exit_node: tensor_lib.Tensor, pfor_ops: List[ops.Operation], fallback_to_while_loop: bool, pfor_config: 'PForConfig'):
    self._fallback_to_while_loop = fallback_to_while_loop
    self._pfor_config = pfor_config
    self._pfor_ops = set(pfor_ops)
    self._pfor_op_ids = set((x._id for x in pfor_ops))
    assert isinstance(exit_node, tensor_lib.Tensor)
    self._while_context = exit_node.op._get_control_flow_context()
    assert isinstance(self._while_context, control_flow_ops.WhileContext)
    self._context_name = self._while_context.name
    self._condition = self._while_context.pivot.op.inputs[0]
    self._is_inside_loop = self.op_is_inside_loop(self._condition.op)
    if self._is_inside_loop:
        for e in self._while_context.loop_exits:
            assert self.op_is_inside_loop(e.op)
    self._exit_switches = []
    self._body_outputs = []
    self._next_iter_control_inputs = []
    self._enter_merges = []
    self._outputs = []
    self._enters = []
    self._direct_enters = []
    for e in self._while_context.loop_exits:
        self._outputs.append(e.op.outputs[0])
        switch = e.op.inputs[0].op
        assert switch.type == 'Switch', switch
        self._exit_switches.append(switch)
        merge = switch.inputs[0].op
        assert merge.type == 'Merge', merge
        self._enter_merges.append(merge)
        enter = merge.inputs[0].op
        assert enter.type == 'Enter', enter
        self._enters.append(enter.outputs[0])
        next_iter = merge.inputs[1].op
        assert next_iter.type == 'NextIteration', next_iter
        self._body_outputs.append(next_iter.inputs[0])
        self._next_iter_control_inputs.append(next_iter.control_inputs)
    self._is_stateful = False
    for op in ops.get_default_graph().get_operations():
        control_flow_context = op._get_control_flow_context()
        if control_flow_context is None:
            continue
        if control_flow_context.name == self._context_name:
            self._is_stateful |= _is_stateful_pfor_op(op)
            if op.type == 'Enter':
                output = op.outputs[0]
                if output not in self._enters:
                    if output.dtype in (dtypes.resource, dtypes.variant):
                        if output not in self._direct_enters:
                            self._direct_enters.append(output)
                    else:
                        self._enters.append(output)

Initializer.

Args:
  exit_node: A tensor output from the while_loop.
  pfor_ops: list of ops inside the current pfor loop.
  fallback_to_while_loop: If True, fallback to while loop when conversion of
    an op is not supported
  pfor_config: PForConfig object used while constructing loop body.


def get_image_features(self, pixel_values: TFModelInputType | None=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, training: bool=False) -> tf.Tensor:
    image_features = self.groupvit.get_image_features(pixel_values=pixel_values, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, training=training)
    return image_features

Returns:
    image_features (`tf.Tensor` of shape `(batch_size, output_dim`): The image embeddings obtained by applying
    the projection layer to the pooled output of [`TFGroupViTVisionModel`].

Examples:

```python
>>> from PIL import Image
>>> import requests
>>> from transformers import AutoProcessor, TFGroupViTModel

>>> model = TFGroupViTModel.from_pretrained("nvidia/groupvit-gcc-yfcc")
>>> processor = AutoProcessor.from_pretrained("nvidia/groupvit-gcc-yfcc")

>>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
>>> image = Image.open(requests.get(url, stream=True).raw)

>>> inputs = processor(images=image, return_tensors="tf")

>>> image_features = model.get_image_features(**inputs)
```


def _warm_start_var_with_vocab(var, current_vocab_path, current_vocab_size, prev_ckpt, prev_vocab_path, previous_vocab_size=-1, current_oov_buckets=0, prev_tensor_name=None, initializer=None, axis=0):
    if not (current_vocab_path and current_vocab_size and prev_ckpt and prev_vocab_path):
        raise ValueError('Invalid args: Must provide all of [current_vocab_path, current_vocab_size, prev_ckpt, prev_vocab_path}.')
    if checkpoint_utils._is_variable(var):
        var = [var]
    elif isinstance(var, list) and all((checkpoint_utils._is_variable(v) for v in var)):
        var = var
    elif isinstance(var, variables_lib.PartitionedVariable):
        var = var._get_variable_list()
    else:
        raise TypeError('var MUST be one of the following: a Variable, list of Variable or PartitionedVariable, but is {}'.format(type(var)))
    if not prev_tensor_name:
        prev_tensor_name = _infer_var_name(var)
    total_v_first_axis = sum((v.get_shape().as_list()[0] for v in var))
    for v in var:
        v_shape = v.get_shape().as_list()
        slice_info = v._get_save_slice_info()
        partition_info = None
        if slice_info:
            partition_info = variable_scope._PartitionInfo(full_shape=slice_info.full_shape, var_offset=slice_info.var_offset)
        if axis == 0:
            new_row_vocab_size = current_vocab_size
            new_col_vocab_size = v_shape[1]
            old_row_vocab_size = previous_vocab_size
            old_row_vocab_file = prev_vocab_path
            new_row_vocab_file = current_vocab_path
            old_col_vocab_file = None
            new_col_vocab_file = None
            num_row_oov_buckets = current_oov_buckets
            num_col_oov_buckets = 0
        elif axis == 1:
            new_row_vocab_size = total_v_first_axis
            new_col_vocab_size = current_vocab_size
            old_row_vocab_size = -1
            old_row_vocab_file = None
            new_row_vocab_file = None
            old_col_vocab_file = prev_vocab_path
            new_col_vocab_file = current_vocab_path
            num_row_oov_buckets = 0
            num_col_oov_buckets = current_oov_buckets
        else:
            raise ValueError('The only supported values for the axis argument are 0 and 1.  Provided axis: {}'.format(axis))
        init = checkpoint_ops._load_and_remap_matrix_initializer(ckpt_path=checkpoint_utils._get_checkpoint_filename(prev_ckpt), old_tensor_name=prev_tensor_name, new_row_vocab_size=new_row_vocab_size, new_col_vocab_size=new_col_vocab_size, old_row_vocab_size=old_row_vocab_size, old_row_vocab_file=old_row_vocab_file, new_row_vocab_file=new_row_vocab_file, old_col_vocab_file=old_col_vocab_file, new_col_vocab_file=new_col_vocab_file, num_row_oov_buckets=num_row_oov_buckets, num_col_oov_buckets=num_col_oov_buckets, initializer=initializer)
        new_init_val = ops.convert_to_tensor(init(shape=v_shape, partition_info=partition_info))
        v._initializer_op = state_ops.assign(v, new_init_val)

Warm-starts given variable from `prev_tensor_name` tensor in `prev_ckpt`.

Use this method when the `var` is backed by vocabulary. This method stitches
the given `var` such that values corresponding to individual features in the
vocabulary remain consistent irrespective of changing order of the features
between old and new vocabularies.

Args:
  var: Current graph's variable that needs to be warm-started (initialized).
    Can be either of the following:
    (i) `Variable`
    (ii) `ResourceVariable`
    (iii) list of `Variable`: The list must contain slices of the same larger
      variable.
    (iv) `PartitionedVariable`
  current_vocab_path: Path to the vocab file used for the given `var`.
  current_vocab_size: An `int` specifying the number of entries in the current
    vocab.
  prev_ckpt: A string specifying the directory with checkpoint file(s) or path
    to checkpoint. The given checkpoint must have tensor with name
    `prev_tensor_name` (if not None) or tensor with name same as given `var`.
  prev_vocab_path: Path to the vocab file used for the tensor in `prev_ckpt`.
  previous_vocab_size: If provided, will constrain previous vocab to the first
    `previous_vocab_size` entries.  -1 means use the entire previous vocab.
  current_oov_buckets: An `int` specifying the number of out-of-vocabulary
    buckets used for given `var`.
  prev_tensor_name: Name of the tensor to lookup in provided `prev_ckpt`. If
    None, we lookup tensor with same name as given `var`.
  initializer: Variable initializer to be used for missing entries.  If None,
    missing entries will be zero-initialized.
  axis: Axis of the variable that the provided vocabulary corresponds to.

Raises:
  ValueError: If required args are not provided.


def _has_valid_tensors(self):
    return self._input_tensors is not None and self._output_tensors

Checks if the input and output tensors have been initialized.

Returns:
  Bool.


def integrate(func: Callable[[types.FloatTensor], types.FloatTensor], lower: types.FloatTensor, upper: types.FloatTensor, method: IntegrationMethod=IntegrationMethod.COMPOSITE_SIMPSONS_RULE, dtype: Optional[tf.DType]=None, name: Optional[str]=None, **kwargs) -> types.FloatTensor:
    with tf.compat.v1.name_scope(name, default_name='integrate', values=[lower, upper]):
        if method == IntegrationMethod.COMPOSITE_SIMPSONS_RULE:
            return simpson(func, lower, upper, dtype=dtype, **kwargs)
        elif method == IntegrationMethod.GAUSS_LEGENDRE:
            return gauss_legendre(func, lower, upper, dtype=dtype, **kwargs)
        else:
            raise ValueError('Unknown method: %s.' % method)

Evaluates definite integral.

#### Example
```python
  f = lambda x: x*x
  a = tf.constant(0.0)
  b = tf.constant(3.0)
  integrate(f, a, b) # 9.0
```

Args:
  func: Represents a function to be integrated. It must be a callable of a
    single `Tensor` parameter and return a `Tensor` of the same shape and
    dtype as its input. It will be called with a `Tesnor` of shape
    `lower.shape + [n]` (where n is integer number of points) and of the same
    `dtype` as `lower`.
  lower: Represents the lower limits of integration. `func` will be integrated
    between each pair of points defined by `lower` and `upper`.
  upper: Same shape and dtype as `lower` representing the upper limits of
    intergation.
  method: Integration method. Instance of IntegrationMethod enum. Default is
    IntegrationMethod.COMPOSITE_SIMPSONS_RULE.
  dtype: Dtype of result. Must be real dtype. Defaults to dtype of `lower`.
  name: The name to give to the ops created by this function.
    Default value: None which maps to 'integrate'.
  **kwargs: Additional parameters for specific integration method.

Returns:
  `Tensor` of the same shape and dtype as `lower`, containing the value of the
  definite integral.

Raises: ValueError if `method` was not recognized.


def build_hlo_module(root: testlib_base.HloInstruction, *instructions: testlib_base.HloInstruction, extra_computations: Sequence[testlib_base.HloComputation] | None=None) -> tuple[testlib_base.HloModule, testlib_base.BufferAssignment]:
    hlo_module = testlib_base.HloModule(root.name())
    hlo_module.add_entry_computation(testlib_base.build_hlo_computation(root, *instructions))
    if extra_computations is not None:
        for computation in extra_computations:
            hlo_module.add_computation(computation)
    return annotate_hlo_module(hlo_module)

Builds an HLO module from a root instruction and its dependencies.

Args:
  root: The root instruction of the module.
  *instructions: The instructions that are dependencies of the root
    instruction.
  extra_computations: Any extra computations that should be added to the
    module.

Returns:
  A tuple containing the HLO module and its buffer assignment.


def take_grad(self, num_required, name=None):
    out = gen_data_flow_ops.resource_accumulator_take_gradient(self._accumulator_ref, num_required, dtype=self._dtype, name=name)
    out.set_shape(self._shape)
    return out

Attempts to extract the average gradient from the accumulator.

The operation blocks until sufficient number of gradients have been
successfully applied to the accumulator.

Once successful, the following actions are also triggered:

- Counter of accumulated gradients is reset to 0.
- Aggregated gradient is reset to 0 tensor.
- Accumulator's internal time step is incremented by 1.

Args:
  num_required: Number of gradients that needs to have been aggregated
  name: Optional name for the operation

Returns:
  A tensor holding the value of the average gradient.

Raises:
  InvalidArgumentError: If num_required < 1


def _GetBcastSubshape(subscripts):
    start = subscripts.find(ellipsis)
    if start == -1:
        return (0, 0)
    remaining = len(subscripts) - (start + len(ellipsis))
    end = -remaining if remaining > 0 else None
    return (start, end)

Returns a tuple denoting the slice mapping to ellipsis.

For a given subscript, returns a tuple (start, end) denoting the start
axis index and the (negative) end axis index respectively. For any input
Tensor `x` described by the subscript, `x[start:end]` would be the slice
represented by the ellipsis. E.g. For `ab...cd` returns `[1, -2]`.

If ellipsis is not present in `subscripts`, returns `(0, 0)`.

Args:
  subscripts: A string denoting the einsum subscript.


class TargetReshaper(TransformerMixin, BaseEstimator):

    def fit(self, y):
        """Fit the transformer to a target y.

        Returns:
            TargetReshaper
                A reference to the current instance of TargetReshaper.
        """
        self.ndim_ = y.ndim
        return self

    def transform(self, y):
        """Makes 1D y 2D.

        Args:
            y : np.ndarray
                Target y to be transformed.

        Returns:
            np.ndarray
                A numpy array, of dimension at least 2.
        """
        if y.ndim == 1:
            return y.reshape(-1, 1)
        return y

    def inverse_transform(self, y):
        """Revert the transformation of transform.

        Args:
            y: np.ndarray
                Transformed numpy array.

        Returns:
            np.ndarray
                If the transformer was fit to a 1D numpy array,
                and a 2D numpy array with a singleton second dimension
                is passed, it will be squeezed back to 1D. Otherwise, it
                will eb left untouched.
        """
        sklearn.base.check_is_fitted(self)
        xp, _ = sklearn.utils._array_api.get_namespace(y)
        if self.ndim_ == 1 and y.ndim == 2:
            return xp.squeeze(y, axis=1)
        return y

Convert 1D targets to 2D and back.

For use in pipelines with transformers that only accept
2D inputs, like OneHotEncoder and OrdinalEncoder.

Attributes:
    ndim_ : int
        Dimensions of y that the transformer was trained on.


def checkout_commit(repo: Repo, commit_id: str):
    current_head = repo.head.commit if repo.head.is_detached else repo.head.ref
    try:
        repo.git.checkout(commit_id)
        yield
    finally:
        repo.git.checkout(current_head)

Context manager that checks out a given commit when entered, but gets back to the reference it was at on exit.
Args:
    repo (`git.Repo`): A git repository (for instance the Transformers repo).
    commit_id (`str`): The commit reference to checkout inside the context manager.


def get_initialized_tpu_systems():
    return _INITIALIZED_TPU_SYSTEMS.copy()

Returns all currently initialized tpu systems.

Returns:
   A dictionary, with tpu name as the key and the tpu topology as the value.


def center_crop(self, image: 'torch.Tensor', size: dict[str, int], **kwargs) -> 'torch.Tensor':
    if size.height is None or size.width is None:
        raise ValueError(f"The size dictionary must have keys 'height' and 'width'. Got {size.keys()}")
    return F.center_crop(image, (size['height'], size['width']))

Center crop an image to `(size["height"], size["width"])`. If the input size is smaller than `crop_size` along
any edge, the image is padded with 0's and then center cropped.

Args:
    image (`"torch.Tensor"`):
        Image to center crop.
    size (`Dict[str, int]`):
        Size of the output image.

Returns:
    `torch.Tensor`: The center cropped image.


def forward(self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor]=None, inputs_embeds: Optional[torch.Tensor]=None, head_mask: Optional[torch.Tensor]=None, output_attentions: bool=False, output_hidden_states: bool=False, return_dict: bool=True):
    if attention_mask is not None:
        attention_mask = invert_mask(attention_mask)
    if input_ids is not None and inputs_embeds is not None:
        raise ValueError('You cannot specify both input_ids and inputs_embeds at the same time')
    elif input_ids is not None:
        inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale
        embed_pos = self.embed_positions(input_ids)
    elif inputs_embeds is not None:
        inputs_embeds = inputs_embeds * self.embed_scale
        position_ids = inputs_embeds[:, :, 0].masked_fill(inputs_embeds[:, :, 0].eq(0), self.embed_positions.padding_idx)
        embed_pos = self.embed_positions(position_ids)
    else:
        raise ValueError('You have to specify either input_ids or inputs_embeds')
    x = inputs_embeds + embed_pos
    x = nn.functional.dropout(x, p=self.dropout, training=self.training)
    x = x.transpose(0, 1)
    encoder_states = () if output_hidden_states else None
    all_attentions = () if output_attentions else None
    if head_mask is not None:
        assert head_mask.size()[0] == len(self.layers), f'The head_mask should be specified for {len(self.layers)} layers, but it is for {head_mask.size()[0]}.'
    for idx, encoder_layer in enumerate(self.layers):
        if output_hidden_states:
            x = x.transpose(0, 1)
            encoder_states += (x,)
            x = x.transpose(0, 1)
        dropout_probability = torch.rand([])
        if self.training and dropout_probability < self.layerdrop:
            attn = None
        else:
            x, attn = encoder_layer(x, attention_mask, layer_head_mask=head_mask[idx] if head_mask is not None else None, output_attentions=output_attentions)
        if output_attentions:
            all_attentions = all_attentions + (attn,)
    x = x.transpose(0, 1)
    if output_hidden_states:
        encoder_states += (x,)
    if not return_dict:
        return tuple((v for v in [x, encoder_states, all_attentions] if v is not None))
    return BaseModelOutput(last_hidden_state=x, hidden_states=encoder_states, attentions=all_attentions)

Args:
    input_ids (`torch.LongTensor`): tokens in the source language of shape
        *(batch, src_len)*
    attention_mask (`torch.LongTensor`): indicating which indices are padding tokens
    inputs_embeds (`torch.FloatTensor`):
        embedding vectors of shape *(batch, src_len, embed_dim)*
    head_mask (`torch.Tensor` of shape `(num_layers, num_heads)`, *optional*):
        Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`:

        - 1 indicates the head is **not masked**,
        - 0 indicates the head is **masked**.

Returns:
    BaseModelOutput or Tuple comprised of:

        - **x** (`torch.Tensor`): the last encoder layer's output of shape *(src_len, batch, embed_dim)*
        - **encoder_states** (`Tuple(torch.FloatTensor)`): all intermediate hidden states of shape *(src_len,
          batch, embed_dim)*. Only populated if *output_hidden_states:* is True.
        - **all_attentions** (`Tuple(torch.FloatTensor)`): Attention weights for each layer.
        During training might not be of length n_layers because of layer dropout.


class FlaxQuestionAnsweringModelOutput(ModelOutput):
    start_logits: Optional[jnp.ndarray] = None
    end_logits: Optional[jnp.ndarray] = None
    hidden_states: Optional[Tuple[jnp.ndarray]] = None
    attentions: Optional[Tuple[jnp.ndarray]] = None

Base class for outputs of question answering models.

Args:
    start_logits (`jnp.ndarray` of shape `(batch_size, sequence_length)`):
        Span-start scores (before SoftMax).
    end_logits (`jnp.ndarray` of shape `(batch_size, sequence_length)`):
        Span-end scores (before SoftMax).
    hidden_states (`tuple(jnp.ndarray)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
        Tuple of `jnp.ndarray` (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(jnp.ndarray)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
        Tuple of `jnp.ndarray` (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.


def experimental_set_type(self, type_proto) -> None:
    with self.graph._c_graph.get() as c_graph:
        if type_proto.type_id not in (full_type_pb2.TFT_UNSET, full_type_pb2.TFT_PRODUCT):
            raise ValueError('error setting the type of ', self.name, ': expected TFT_UNSET or TFT_PRODUCT, got ', type_proto.type_id)
        with c_api_util.tf_buffer(type_proto.SerializeToString()) as serialized:
            pywrap_tf_session.SetFullType(c_graph, self._c_op, serialized)

Sets the corresponding node's `experimental_type` field.

See the description of `NodeDef.experimental_type` for more info.

Args:
  type_proto: A FullTypeDef proto message. The root type_if of this object
    must be `TFT_PRODUCT`, even for ops which only have a singlre return
    value.


def __call__(self, *args, **kwargs):
    return super().__call__(*args, **kwargs)

Translate the text(s) given as inputs.

Args:
    args (`str` or `List[str]`):
        Texts to be translated.
    return_tensors (`bool`, *optional*, defaults to `False`):
        Whether or not to include the tensors of predictions (as token indices) in the outputs.
    return_text (`bool`, *optional*, defaults to `True`):
        Whether or not to include the decoded texts in the outputs.
    clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`):
        Whether or not to clean up the potential extra spaces in the text output.
    src_lang (`str`, *optional*):
        The language of the input. Might be required for multilingual models. Will not have any effect for
        single pair translation models
    tgt_lang (`str`, *optional*):
        The language of the desired output. Might be required for multilingual models. Will not have any effect
        for single pair translation models
    generate_kwargs:
        Additional keyword arguments to pass along to the generate method of the model (see the generate method
        corresponding to your framework [here](./text_generation)).

Return:
    A list or a list of list of `dict`: Each result comes as a dictionary with the following keys:

    - **translation_text** (`str`, present when `return_text=True`) -- The translation.
    - **translation_token_ids** (`torch.Tensor` or `tf.Tensor`, present when `return_tensors=True`) -- The
      token ids of the translation.


def load_data():
    dirname = os.path.join('datasets', 'fashion-mnist')
    base = 'https://storage.googleapis.com/tensorflow/tf-keras-datasets/'
    files = ['train-labels-idx1-ubyte.gz', 'train-images-idx3-ubyte.gz', 't10k-labels-idx1-ubyte.gz', 't10k-images-idx3-ubyte.gz']
    paths = []
    for fname in files:
        paths.append(get_file(fname, origin=base + fname, cache_subdir=dirname))
    with gzip.open(paths[0], 'rb') as lbpath:
        y_train = np.frombuffer(lbpath.read(), np.uint8, offset=8)
    with gzip.open(paths[1], 'rb') as imgpath:
        x_train = np.frombuffer(imgpath.read(), np.uint8, offset=16).reshape(len(y_train), 28, 28)
    with gzip.open(paths[2], 'rb') as lbpath:
        y_test = np.frombuffer(lbpath.read(), np.uint8, offset=8)
    with gzip.open(paths[3], 'rb') as imgpath:
        x_test = np.frombuffer(imgpath.read(), np.uint8, offset=16).reshape(len(y_test), 28, 28)
    return ((x_train, y_train), (x_test, y_test))

Loads the Fashion-MNIST dataset.

This is a dataset of 60,000 28x28 grayscale images of 10 fashion categories,
along with a test set of 10,000 images. This dataset can be used as
a drop-in replacement for MNIST.

The classes are:

| Label | Description |
|:-----:|-------------|
|   0   | T-shirt/top |
|   1   | Trouser     |
|   2   | Pullover    |
|   3   | Dress       |
|   4   | Coat        |
|   5   | Sandal      |
|   6   | Shirt       |
|   7   | Sneaker     |
|   8   | Bag         |
|   9   | Ankle boot  |

Returns:

Tuple of NumPy arrays: `(x_train, y_train), (x_test, y_test)`.

**`x_train`**: `uint8` NumPy array of grayscale image data with shapes
  `(60000, 28, 28)`, containing the training data.

**`y_train`**: `uint8` NumPy array of labels (integers in range 0-9)
  with shape `(60000,)` for the training data.

**`x_test`**: `uint8` NumPy array of grayscale image data with shapes
  (10000, 28, 28), containing the test data.

**`y_test`**: `uint8` NumPy array of labels (integers in range 0-9)
  with shape `(10000,)` for the test data.

Example:

```python
(x_train, y_train), (x_test, y_test) = fashion_mnist.load_data()
assert x_train.shape == (60000, 28, 28)
assert x_test.shape == (10000, 28, 28)
assert y_train.shape == (60000,)
assert y_test.shape == (10000,)
```

License:

The copyright for Fashion-MNIST is held by Zalando SE.
Fashion-MNIST is licensed under the [MIT license](
    https://github.com/zalandoresearch/fashion-mnist/blob/master/LICENSE).


def from_fields_and_rank(cls, fields: Mapping[str, _FieldValue], rank: int, validate: bool=False, dtype: Optional[dtypes.DType]=None) -> 'StructuredTensor':
    if not fields:
        raise ValueError('Must provide at least one field')
    if not isinstance(rank, int):
        raise ValueError('rank must be an integer')
    if rank < 0:
        raise ValueError('rank must be nonnegative')
    fields = {k: _convert_to_structured_field_value(v) for k, v in fields.items()}
    if dtype is None:
        dtype = _find_shape_dtype(fields, None, None)
    fields = _fields_with_dtype(fields, dtype)
    shape = _shape_from_fields(fields, rank, dtype)
    if rank > 1:
        shape = shape._with_num_row_partitions(rank - 1)
    new_rp = shape._row_partitions
    fields = {k: _replace_row_partitions(v, new_rp) for k, v in fields.items()}
    return StructuredTensor(fields=fields, ragged_shape=shape)

Creates a `StructuredTensor` from a nonempty dictionary of fields.

Note that if the shape dtype is not specified, the shape dtype will be
inferred from any fields that have a shape dtype. If fields differ, then
int64 will be preferred to int32, because coercing from int32 to int64 is
safer than coercing from int64 to int32.

If there are no ragged fields, then it will be int64 by default, but this
will be changed to int32 in the future.

Args:
  fields: A dictionary mapping from string to `Tensor`, `RaggedTensor`, or
    `StructuredTensor`, providing the values for individual fields in each
    structure.  If `rank > 0`, then every tensor in `fields` must have the
    same shape in the first `rank` dimensions. Cannot be empty.
  rank: The rank of the resulting structured tensor.
  validate: If true, then add runtime validation ops that check that the
    field values all have compatible shapes in the outer `rank` dimensions.
  dtype: If specified, then forces dtype of the shape to be this.

Returns:
  A `StructuredTensor`.
Examples:
  >>> tf.experimental.StructuredTensor.from_fields_and_rank(
  ...     {'x': 1, 'y': [1, 2, 3]}, 0)
  <StructuredTensor(
    fields={
      "x": tf.Tensor(1, shape=(), dtype=int32),
      "y": tf.Tensor([1 2 3], shape=(3,), dtype=int32)},
    shape=())>
  >>> StructuredTensor.from_fields_and_rank({'foo': [1, 2], 'bar': [3, 4]},
  ...                              1)
  <StructuredTensor(
    fields={
      "bar": tf.Tensor([3 4], shape=(2,), dtype=int32),
      "foo": tf.Tensor([1 2], shape=(2,), dtype=int32)},
    shape=(2,))>


def rgb_to_hsv(images, data_format=None):
    if any_symbolic_tensors((images,)):
        return RGBToHSV(data_format=data_format).symbolic_call(images)
    return backend.image.rgb_to_hsv(images, data_format=data_format)

Convert RGB images to HSV.

`images` must be of float dtype, and the output is only well defined if the
values in `images` are in `[0, 1]`.

All HSV values are in `[0, 1]`. A hue of `0` corresponds to pure red, `1/3`
is pure green, and `2/3` is pure blue.

Args:
    images: Input image or batch of images. Must be 3D or 4D.
    data_format: A string specifying the data format of the input tensor.
        It can be either `"channels_last"` or `"channels_first"`.
        `"channels_last"` corresponds to inputs with shape
        `(batch, height, width, channels)`, while `"channels_first"`
        corresponds to inputs with shape `(batch, channels, height, width)`.
        If not specified, the value will default to
        `keras.config.image_data_format`.

Returns:
    HSV image or batch of HSV images.

Examples:

>>> import numpy as np
>>> from keras import ops
>>> x = np.random.random((2, 4, 4, 3))
>>> y = ops.image.rgb_to_hsv(x)
>>> y.shape
(2, 4, 4, 3)

>>> x = np.random.random((4, 4, 3)) # Single RGB image
>>> y = ops.image.rgb_to_hsv(x)
>>> y.shape
(4, 4, 3)

>>> x = np.random.random((2, 3, 4, 4))
>>> y = ops.image.rgb_to_hsv(x, data_format="channels_first")
>>> y.shape
(2, 3, 4, 4)


def forward(self, hidden_states: torch.FloatTensor, attention_mask: torch.FloatTensor, layer_head_mask: torch.FloatTensor, output_attentions: Optional[bool]=False) -> Tuple[torch.FloatTensor, Optional[torch.FloatTensor]]:
    residual = hidden_states
    hidden_states, attn_weights, _ = self.self_attn(hidden_states=hidden_states, attention_mask=attention_mask, layer_head_mask=layer_head_mask, output_attentions=output_attentions)
    hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
    hidden_states = residual + hidden_states
    hidden_states = self.self_attn_layer_norm(hidden_states)
    residual = hidden_states
    hidden_states = self.activation_fn(self.fc1(hidden_states))
    hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
    hidden_states = self.fc2(hidden_states)
    hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
    hidden_states = residual + hidden_states
    hidden_states = self.final_layer_norm(hidden_states)
    if hidden_states.dtype == torch.float16 and (torch.isinf(hidden_states).any() or torch.isnan(hidden_states).any()):
        clamp_value = torch.finfo(hidden_states.dtype).max - 1000
        hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
    outputs = (hidden_states,)
    if output_attentions:
        outputs += (attn_weights,)
    return outputs

Args:
    hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
    attention_mask (`torch.FloatTensor`): attention mask of size
        `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
    layer_head_mask (`torch.FloatTensor`): mask for attention heads in a given layer of size
        `(encoder_attention_heads,)`.
    output_attentions (`bool`, *optional*):
        Whether or not to return the attentions tensors of all attention layers. See `attentions` under
        returned tensors for more detail.


def seed(self, initial_seed: Union[None, tf.Tensor]):
    if initial_seed is None:
        self._update_seed.assign(True)
    else:
        self._update_seed.assign(False)
    self._seed.assign(tfp.random.sanitize_seed(initial_seed))

Sets a new value of the random seed.

Args:
  initial_seed: see `self.seed` for details.


class SqlChain:
    nodes: Dict[str, SqlNode] = None
    root: Optional[SqlNode] = None
    current: Optional[SqlNode] = None
    user_pipeline: Optional[beam.Pipeline] = None

    def __post_init__(self):
        if not self.nodes:
            self.nodes = {}

    @progress_indicated
    def to_pipeline(self) -> beam.Pipeline:
        """Converts the chain into a beam pipeline."""
        pipeline_to_execute = self.root.to_pipeline(self.user_pipeline)
        pipeline_to_execute.contains_external_transforms = True
        return pipeline_to_execute

    def append(self, node: SqlNode) -> 'SqlChain':
        """Appends a node to the chain."""
        if self.current:
            self.current.next = node
        else:
            self.root = node
        self.current = node
        self.nodes[node.output_name] = node
        return self

    def get(self, output_name: str) -> Optional[SqlNode]:
        """Gets a node from the chain based on the given output_name."""
        return self.nodes.get(output_name, None)

A chain of SqlNodes.

Attributes:
  nodes: all nodes by their output_names.
  root: the first SqlNode applied chronologically.
  current: the last node applied.
  user_pipeline: the user defined pipeline this chain originates from. If
    None, the whole chain just computes from raw values in queries.
    Otherwise, at least some of the nodes in chain has queried against
    PCollections.


class Speech2TextTokenizer(PreTrainedTokenizer):
    vocab_files_names = VOCAB_FILES_NAMES
    model_input_names = ['input_ids', 'attention_mask']
    prefix_tokens: List[int] = []

    def __init__(self, vocab_file, spm_file, bos_token='<s>', eos_token='</s>', pad_token='<pad>', unk_token='<unk>', do_upper_case=False, do_lower_case=False, tgt_lang=None, lang_codes=None, additional_special_tokens=None, sp_model_kwargs: Optional[Dict[str, Any]]=None, **kwargs) -> None:
        self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs
        self.do_upper_case = do_upper_case
        self.do_lower_case = do_lower_case
        self.encoder = load_json(vocab_file)
        self.decoder = {v: k for k, v in self.encoder.items()}
        self.spm_file = spm_file
        self.sp_model = load_spm(spm_file, self.sp_model_kwargs)
        if lang_codes is not None:
            self.lang_codes = lang_codes
            self.langs = LANGUAGES[lang_codes]
            self.lang_tokens = [f'<lang:{lang}>' for lang in self.langs]
            self.lang_code_to_id = {lang: self.sp_model.PieceToId(f'<lang:{lang}>') for lang in self.langs}
            if additional_special_tokens is not None:
                additional_special_tokens = self.lang_tokens + additional_special_tokens
            else:
                additional_special_tokens = self.lang_tokens
            self._tgt_lang = tgt_lang if tgt_lang is not None else self.langs[0]
            self.set_tgt_lang_special_tokens(self._tgt_lang)
        else:
            self.lang_code_to_id = {}
        super().__init__(bos_token=bos_token, eos_token=eos_token, unk_token=unk_token, pad_token=pad_token, do_upper_case=do_upper_case, do_lower_case=do_lower_case, tgt_lang=tgt_lang, lang_codes=lang_codes, sp_model_kwargs=self.sp_model_kwargs, additional_special_tokens=additional_special_tokens, **kwargs)

    @property
    def vocab_size(self) -> int:
        return len(self.encoder)

    def get_vocab(self) -> Dict:
        vocab = self.encoder.copy()
        vocab.update(self.added_tokens_encoder)
        return vocab

    @property
    def tgt_lang(self) -> str:
        return self._tgt_lang

    @tgt_lang.setter
    def tgt_lang(self, new_tgt_lang) -> None:
        self._tgt_lang = new_tgt_lang
        self.set_tgt_lang_special_tokens(new_tgt_lang)

    def set_tgt_lang_special_tokens(self, tgt_lang: str) -> None:
        """Reset the special tokens to the target language setting. prefix=[eos, tgt_lang_code] and suffix=[eos]."""
        lang_code_id = self.lang_code_to_id[tgt_lang]
        self.prefix_tokens = [lang_code_id]

    def _tokenize(self, text: str) -> List[str]:
        return self.sp_model.encode(text, out_type=str)

    def _convert_token_to_id(self, token):
        return self.encoder.get(token, self.encoder[self.unk_token])

    def _convert_id_to_token(self, index: int) -> str:
        """Converts an index (integer) in a token (str) using the decoder."""
        return self.decoder.get(index, self.unk_token)

    def convert_tokens_to_string(self, tokens: List[str]) -> str:
        """Converts a sequence of tokens (strings for sub-words) in a single string."""
        current_sub_tokens = []
        out_string = ''
        for token in tokens:
            if token in self.all_special_tokens:
                decoded = self.sp_model.decode(current_sub_tokens)
                out_string += (decoded.upper() if self.do_upper_case else decoded) + token + ' '
                current_sub_tokens = []
            else:
                current_sub_tokens.append(token)
        decoded = self.sp_model.decode(current_sub_tokens)
        out_string += decoded.upper() if self.do_upper_case else decoded
        return out_string.strip()

    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None) -> List[int]:
        """Build model inputs from a sequence by appending eos_token_id."""
        if token_ids_1 is None:
            return self.prefix_tokens + token_ids_0 + [self.eos_token_id]
        return self.prefix_tokens + token_ids_0 + token_ids_1 + [self.eos_token_id]

    def get_special_tokens_mask(self, token_ids_0: List[int], token_ids_1: Optional[List[int]]=None, already_has_special_tokens: bool=False) -> List[int]:
        """
        Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
        special tokens using the tokenizer `prepare_for_model` method.

        Args:
            token_ids_0 (`List[int]`):
                List of IDs.
            token_ids_1 (`List[int]`, *optional*):
                Optional second list of IDs for sequence pairs.
            already_has_special_tokens (`bool`, *optional*, defaults to `False`):
                Whether or not the token list is already formatted with special tokens for the model.

        Returns:
            `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
        """
        if already_has_special_tokens:
            return super().get_special_tokens_mask(token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True)
        prefix_ones = [1] * len(self.prefix_tokens)
        suffix_ones = [1]
        if token_ids_1 is None:
            return prefix_ones + [0] * len(token_ids_0) + suffix_ones
        return prefix_ones + [0] * len(token_ids_0) + [0] * len(token_ids_1) + suffix_ones

    def __getstate__(self) -> Dict:
        state = self.__dict__.copy()
        state['sp_model'] = None
        return state

    def __setstate__(self, d: Dict) -> None:
        self.__dict__ = d
        if not hasattr(self, 'sp_model_kwargs'):
            self.sp_model_kwargs = {}
        self.sp_model = load_spm(self.spm_file, self.sp_model_kwargs)

    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str]=None) -> Tuple[str]:
        save_dir = Path(save_directory)
        assert save_dir.is_dir(), f'{save_directory} should be a directory'
        vocab_save_path = save_dir / ((filename_prefix + '-' if filename_prefix else '') + self.vocab_files_names['vocab_file'])
        spm_save_path = save_dir / ((filename_prefix + '-' if filename_prefix else '') + self.vocab_files_names['spm_file'])
        save_json(self.encoder, vocab_save_path)
        if os.path.abspath(self.spm_file) != os.path.abspath(spm_save_path) and os.path.isfile(self.spm_file):
            copyfile(self.spm_file, spm_save_path)
        elif not os.path.isfile(self.spm_file):
            with open(spm_save_path, 'wb') as fi:
                content_spiece_model = self.sp_model.serialized_model_proto()
                fi.write(content_spiece_model)
        return (str(vocab_save_path), str(spm_save_path))

Construct an Speech2Text tokenizer.

This tokenizer inherits from [`PreTrainedTokenizer`] which contains some of the main methods. Users should refer to
the superclass for more information regarding such methods.

Args:
    vocab_file (`str`):
        File containing the vocabulary.
    spm_file (`str`):
        Path to the [SentencePiece](https://github.com/google/sentencepiece) model file
    bos_token (`str`, *optional*, defaults to `"<s>"`):
        The beginning of sentence token.
    eos_token (`str`, *optional*, defaults to `"</s>"`):
        The end of sentence token.
    unk_token (`str`, *optional*, defaults to `"<unk>"`):
        The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
        token instead.
    pad_token (`str`, *optional*, defaults to `"<pad>"`):
        The token used for padding, for example when batching sequences of different lengths.
    do_upper_case (`bool`, *optional*, defaults to `False`):
       Whether or not to uppercase the output when decoding.
    do_lower_case (`bool`, *optional*, defaults to `False`):
        Whether or not to lowercase the input when tokenizing.
    tgt_lang (`str`, *optional*):
        A string representing the target language.
    sp_model_kwargs (`dict`, *optional*):
        Will be passed to the `SentencePieceProcessor.__init__()` method. The [Python wrapper for
        SentencePiece](https://github.com/google/sentencepiece/tree/master/python) can be used, among other things,
        to set:

        - `enable_sampling`: Enable subword regularization.
        - `nbest_size`: Sampling parameters for unigram. Invalid for BPE-Dropout.

          - `nbest_size = {0,1}`: No sampling is performed.
          - `nbest_size > 1`: samples from the nbest_size results.
          - `nbest_size < 0`: assuming that nbest_size is infinite and samples from the all hypothesis (lattice)
            using forward-filtering-and-backward-sampling algorithm.

        - `alpha`: Smoothing parameter for unigram sampling, and dropout probability of merge operations for
          BPE-dropout.

    **kwargs
        Additional keyword arguments passed along to [`PreTrainedTokenizer`]