def get_text_features(self, input_ids: Optional[torch.Tensor]=None, attention_mask: Optional[torch.Tensor]=None, position_ids: Optional[torch.Tensor]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None) -> torch.FloatTensor:
    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
    output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    text_outputs = self.text_model(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
    pooled_output = text_outputs[1] if return_dict is not None else text_outputs.pooler_output
    text_features = self.text_projection(pooled_output)
    text_features = F.normalize(text_features, dim=-1)
    return text_features

Returns:
    text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by
    applying the projection layer to the pooled output of [`ClapTextModel`].

Examples:

```python
>>> from transformers import AutoTokenizer, ClapModel

>>> model = ClapModel.from_pretrained("laion/clap-htsat-unfused")
>>> tokenizer = AutoTokenizer.from_pretrained("laion/clap-htsat-unfused")

>>> inputs = tokenizer(["the sound of a cat", "the sound of a dog"], padding=True, return_tensors="pt")
>>> text_features = model.get_text_features(**inputs)
```


def sparse_add_v2(a, b, threshold=0):
    sparse_classes = (sparse_tensor.SparseTensor, sparse_tensor.SparseTensorValue)
    if not any((isinstance(inp, sparse_classes) for inp in [a, b])):
        raise TypeError('At least one input should be SparseTensor; do you mean to use tf.add()?')
    if all((isinstance(inp, sparse_classes) for inp in [a, b])):
        a = _convert_to_sparse_tensor(a)
        b = _convert_to_sparse_tensor(b)
        threshold = ops.convert_to_tensor(threshold, dtype=a.values.dtype.real_dtype.base_dtype, name='threshold')
        output_ind, output_val, output_shape = gen_sparse_ops.sparse_add(a.indices, a.values, a.dense_shape, b.indices, b.values, b.dense_shape, threshold)
        a.get_shape().assert_is_compatible_with(b.get_shape())
        static_shape = array_ops.broadcast_static_shape(a.get_shape(), b.get_shape())
        if static_shape.is_fully_defined():
            output_shape = static_shape.as_list()
        return sparse_tensor.SparseTensor(output_ind, output_val, output_shape)
    else:
        if isinstance(b, sparse_classes):
            a, b = (b, a)
        return gen_sparse_ops.sparse_tensor_dense_add(a.indices, a.values, a.dense_shape, b)

Adds two tensors, at least one of each is a `SparseTensor`.

If one `SparseTensor` and one `Tensor` are passed in, returns a `Tensor`.  If
both arguments are `SparseTensor`s, this returns a `SparseTensor`.  The order
of arguments does not matter.  Use vanilla `tf.add()` for adding two dense
`Tensor`s.

The shapes of the two operands must match: broadcasting is not supported.

The indices of any input `SparseTensor` are assumed ordered in standard
lexicographic order.  If this is not the case, before this step run
`SparseReorder` to restore index ordering.

If both arguments are sparse, we perform "clipping" as follows.  By default,
if two values sum to zero at some index, the output `SparseTensor` would still
include that particular location in its index, storing a zero in the
corresponding value slot.  To override this, callers can specify `threshold`,
indicating that if the sum has a magnitude strictly smaller than `threshold`,
its corresponding value and index would then not be included.  In particular,
`threshold == 0.0` (default) means everything is kept and actual thresholding
happens only for a positive value.

For example, suppose the logical sum of two sparse operands is (densified):

    [       2]
    [.1     0]
    [ 6   -.2]

Then,

* `threshold == 0` (the default): all 5 index/value pairs will be
    returned.
* `threshold == 0.11`: only .1 and 0 will vanish, and the remaining three
    index/value pairs will be returned.
* `threshold == 0.21`: .1, 0, and -.2 will vanish.

Args:
  a: The first operand; `SparseTensor` or `Tensor`.
  b: The second operand; `SparseTensor` or `Tensor`. At least one operand
    must be sparse.
  threshold: A 0-D `Tensor`. The magnitude threshold that determines if an
    output value/index pair takes space. Its dtype should match that of the
    values if they are real; if the latter are complex64/complex128, then the
    dtype should be float32/float64, correspondingly.

Returns:
  A `SparseTensor` or a `Tensor`, representing the sum.

Raises:
  TypeError: If both `a` and `b` are `Tensor`s.  Use `tf.add()` instead.


class VarianceScaling(Initializer):

    def __init__(self, scale=1.0, mode='fan_in', distribution='truncated_normal', seed=None):
        if scale <= 0.0:
            raise ValueError(f'Argument `scale` must be a positive float. Received: {scale}')
        if mode not in {'fan_in', 'fan_out', 'fan_avg'}:
            raise ValueError(f"Argument `mode` should be one of ('fan_in', 'fan_out', 'fan_avg'). Received: {mode}")
        distribution = distribution.lower()
        if distribution == 'normal':
            distribution = 'truncated_normal'
        if distribution not in {'uniform', 'truncated_normal', 'untruncated_normal'}:
            raise ValueError(f"Argument `distribution` should be one of ('uniform', 'truncated_normal', 'untruncated_normal'). Received: {distribution}")
        self.scale = scale
        self.mode = mode
        self.distribution = distribution
        self.seed = seed
        self._random_generator = _RandomGenerator(seed)

    def __call__(self, shape, dtype=dtypes.float32, **kwargs):
        """Returns a tensor object initialized as specified by the initializer.

    Args:
      shape: Shape of the tensor.
      dtype: Optional dtype of the tensor. Only floating point types are
        supported.
      **kwargs: Additional keyword arguments.

    Raises:
      ValueError: If the dtype is not floating point
    """
        self._validate_kwargs(kwargs)
        dtype = _assert_float_dtype(dtype)
        scale = self.scale
        fan_in, fan_out = _compute_fans(shape)
        if _PARTITION_SHAPE in kwargs:
            shape = kwargs[_PARTITION_SHAPE]
        if self.mode == 'fan_in':
            scale /= max(1.0, fan_in)
        elif self.mode == 'fan_out':
            scale /= max(1.0, fan_out)
        else:
            scale /= max(1.0, (fan_in + fan_out) / 2.0)
        if self.distribution == 'truncated_normal':
            stddev = math.sqrt(scale) / 0.8796256610342398
            return self._random_generator.truncated_normal(shape, 0.0, stddev, dtype)
        elif self.distribution == 'untruncated_normal':
            stddev = math.sqrt(scale)
            return self._random_generator.random_normal(shape, 0.0, stddev, dtype)
        else:
            limit = math.sqrt(3.0 * scale)
            return self._random_generator.random_uniform(shape, -limit, limit, dtype)

    def get_config(self):
        return {'scale': self.scale, 'mode': self.mode, 'distribution': self.distribution, 'seed': self.seed}

Initializer capable of adapting its scale to the shape of weights tensors.

Initializers allow you to pre-specify an initialization strategy, encoded in
the Initializer object, without knowing the shape and dtype of the variable
being initialized.

With `distribution="truncated_normal" or "untruncated_normal"`, samples are
drawn from a truncated/untruncated normal distribution with a mean of zero and
a standard deviation (after truncation, if used) `stddev = sqrt(scale / n)`
where n is:

  - number of input units in the weight tensor, if mode = "fan_in"
  - number of output units, if mode = "fan_out"
  - average of the numbers of input and output units, if mode = "fan_avg"

With `distribution="uniform"`, samples are drawn from a uniform distribution
within [-limit, limit], with `limit = sqrt(3 * scale / n)`.

Examples:

>>> def make_variables(k, initializer):
...   return (tf.Variable(initializer(shape=[k], dtype=tf.float32)),
...           tf.Variable(initializer(shape=[k, k], dtype=tf.float32)))
>>> v1, v2 = make_variables(3, tf.initializers.VarianceScaling(scale=1.))
>>> v1
<tf.Variable ... shape=(3,) ... numpy=array([...], dtype=float32)>
>>> v2
<tf.Variable ... shape=(3, 3) ... numpy=
...
>>> make_variables(4, tf.initializers.VarianceScaling(distribution='uniform'))
(<tf.Variable...shape=(4,) dtype=float32...>, <tf.Variable...shape=(4, 4) ...

Args:
  scale: Scaling factor (positive float).
  mode: One of "fan_in", "fan_out", "fan_avg".
  distribution: Random distribution to use. One of "truncated_normal",
    "untruncated_normal" and  "uniform".
  seed: A Python integer. Used to create random seeds. See
    `tf.random.set_seed` for behavior.

Raises:
  ValueError: In case of an invalid value for the "scale", mode" or
    "distribution" arguments.


def img_to_array(img, data_format=None, dtype=None):
    data_format = backend.standardize_data_format(data_format)
    if dtype is None:
        dtype = backend.floatx()
    x = np.asarray(img, dtype=dtype)
    if len(x.shape) == 3:
        if data_format == 'channels_first':
            x = x.transpose(2, 0, 1)
    elif len(x.shape) == 2:
        if data_format == 'channels_first':
            x = x.reshape((1, x.shape[0], x.shape[1]))
        else:
            x = x.reshape((x.shape[0], x.shape[1], 1))
    else:
        raise ValueError(f'Unsupported image shape: {x.shape}')
    return x

Converts a PIL Image instance to a NumPy array.

Example:

```python
from PIL import Image
img_data = np.random.random(size=(100, 100, 3))
img = keras.utils.array_to_img(img_data)
array = keras.utils.image.img_to_array(img)
```

Args:
    img: Input PIL Image instance.
    data_format: Image data format, can be either `"channels_first"` or
        `"channels_last"`. Defaults to `None`, in which case the global
        setting `keras.backend.image_data_format()` is used (unless you
        changed it, it defaults to `"channels_last"`).
    dtype: Dtype to use. `None` means the global setting
        `keras.backend.floatx()` is used (unless you changed it, it
        defaults to `"float32"`).

Returns:
    A 3D NumPy array.


def __init__(self, *, start_date, end_date, holiday_calendar, backward=False):
    self._start_date = start_date
    self._end_date = end_date
    self._holiday_calendar = holiday_calendar
    self._backward = backward

Initializes the schedule.

Initializes a schedule with a given date range and holiday calendar.

The schedule includes all business days between and including `start_date`
and `end_date`.

Can create multiple schedules simultaneously. The start and end dates may
have any (compatible) shape. The `DateTensor` returned by `dates()` has the
shape `start_date.shape + (n,)`, where `n` is the maximum length of
schedules in the batch. If schedules have different lengths, the extra
elements will be padded with extra `end_date` elements at the end, if
`backward=False` and with extra `start_date` elements in the beginning if
`backward=True`. In all cases each schedule in the batch is monotonic.

#### Example Usage (Non-batch)

```python
  start_date = datetime.dates_from_tuples([(2020, 3, 19)])
  end_date = datetime.dates_from_tuples([(2021, 3, 25)])
  holiday_calendar = datetime.HolidayCalendar(start_year=2020,
                                              end_year=2021)
  schedule = datetime.BusinessDaysSchedule(
      start_date=start_date,
      end_date=end_date,
      holiday_calendar=holiday_calendar,
      roll_convention=datetime.BusinessDayConvention.FOLLOWING,
      backward=False).dates()
  # schedule is a DateTensor of
  # [[(2020, 3, 19), (2020, 3, 20), (2020, 3, 23), (2020, 3, 24),
  #   (2021, 3, 25)]] regardless of `backward`.
```

#### Example Usage (Batch)

```python
  start_date = datetime.dates_from_tuples([(2020, 3, 19), (2020, 4, 15)])
  end_date = datetime.dates_from_tuples([(2021, 3, 13), (2021, 3, 17)])
  schedule = datetime.BusinessDaysSchedule(
      start_dates,
      end_dates,
      datetime.HolidayCalendar(start_year=2020, end_year=2021),
      backward=False).dates()
  # Returns DateTensor of
  # [[(2020, 3, 19), (2020, 3, 20), (2020, 3, 23), (2020, 3, 24),
  #   (2021, 3, 25)],
  # [(2020, 3, 13), (2020, 3, 16), (2020, 3, 17), (2020, 3, 17),
  #  (2021, 3, 17)]], if `backward` is True.
  # [[(2020, 3, 19), (2020, 3, 20), (2020, 3, 23), (2020, 3, 24),
  #   (2021, 3, 25)],
  # [(2020, 3, 13), (2020, 3, 13), (2020, 3, 13), (2020, 3, 16),
  #  (2021, 3, 17)]], if `backward` is True.
```
Args:
  start_date: `dates.DateTensor`. Defines the lower boundary of schedule. If
    `backward=True` must be broadcastable to `end_date`, otherwise has
    arbitrary shape.
  end_date: `dates.DateTensor`. Defines the upper boundary of the schedule.
    If `backward=False` must be broadcastable to `start_date`, otherwise has
    arbitrary shape.
  holiday_calendar: `dates.HolidayCalendar` that defines which days will be
    included.
  backward: Python `bool`. Defines the way padding is applied in case of
    batching. If schedules in a batch have different lengths, the extra
    elements will be padded with extra `end_date` elements at the end, if
    `backward=False` and with extra `start_date` elements in the beginning
    if `backward=True`.


def assertOutputStateMatches(self, **has_output):
    output_types = {'stdout', 'stderr', 'returncode'}
    assert len(output_types) == len(has_output)
    for output_type in output_types:
        output_value = getattr(self, output_type)
        if has_output[output_type]:
            self.assertTrue(output_value, output_type + ' unexpectedly empty')
        else:
            value = str(output_value)
            if len(value) > 50:
                value = value[:47] + '...'
            self.assertFalse(output_value, f'Unexpected output to {output_type}: {value!r}')

Check that the output state matches expectations.

If, for example, you expect the program to print something to stdout and
nothing to stderr before exiting with an error code, you would write
assertOutputStateMatches(stdout=True, stderr=False, returncode=True).

Args:
  **has_output: Whether each output type should have output.


def on_execution(self, execution_index, execution):

Monitor method for top-level execution events.

Return values (if any) are ignored by the associated DebugDataReader.

Args:
  execution_index: The index of the top-level execution event, as an int.
  execution: An Execution data object, for a top-level op or function
    execution event.


def _NthElementGrad(op: ops.Operation, grad):
    input = op.inputs[0]
    output = op.outputs[0]
    indicators = math_ops.cast(math_ops.equal(array_ops.expand_dims(output, -1), input), grad.dtype)
    grad = array_ops.expand_dims(grad, -1)
    num_selected = array_ops.expand_dims(math_ops.reduce_sum(indicators, -1), -1)
    return [math_ops.divide(indicators, num_selected) * grad, None]

Return the gradients for NthElement.

Args:
  op: The NthElementOp for which we need to generate gradients.
  grad: Tensor. The gradients passed to the NthElementOp

Returns:
  A list of two tensors, the first being the gradient w.r.t. the input,
  the second being the gradient w.r.t. the N (None).


def get_current_remat_mode():
    remat_scope_stack = global_state.get_global_attribute('remat_scope_stack')
    if not remat_scope_stack:
        return None
    active_scope = remat_scope_stack[-1]
    return RematMode(active_scope.mode, active_scope.output_size_threshold, active_scope.layer_names)

Get the current rematerialization mode and associated settings.

Returns:
    RematMode or None: The current rematerialization mode, or None if not
    set.


class _Entrypoint:
    module: str
    name: str
    exported_symbol: exported_api.ExportedSymbol

    def get_import(self, file_prefixes_to_strip: Sequence[str], module_prefix: str, use_lazy_loading: bool) -> str:
        """Returns the import statement for this entrypoint.

    Args:
      file_prefixes_to_strip: List of prefixes to strip from the file name.
      module_prefix: A prefix to add to the import.
      use_lazy_loading: Whether to use lazy loading or not.
    """
        module_import_path = _get_import_path(self.exported_symbol.file_name, file_prefixes_to_strip, module_prefix)
        alias = ''
        symbol_name = self.exported_symbol.symbol_name
        if self.name != symbol_name:
            alias = f' as {self.name}'
        if not use_lazy_loading:
            return f'from {module_import_path} import {symbol_name}{alias} # line: {self.exported_symbol.line_no}'
        else:
            return f"  '{self.name}': ('{module_import_path}', '{symbol_name}'), # line: {self.exported_symbol.line_no}"

An entrypoint that was exposed by the use of a decorator.

Attributes:
  module: The public module that the symbol was exposed to. For example:
    tensorflow.io.
  name: The name the symbol was exported as. For example: decode_png.
  exported_symbol: The symbol that this entrypoint refers back to.


def placeholder_value(self, placeholder_context) -> Any:

Creates a placeholder for tracing.

tf.funcion traces with the placeholder value rather than the actual value.
For example, a placeholder value can represent multiple different
actual values. This means that the trace generated with that placeholder
value is more general and reusable which saves expensive retracing.

Args:
  placeholder_context: A context reserved for internal/future usage.
For the `Fruit` example shared above, implementing:

```python
class FruitTraceType:
  def placeholder_value(self, placeholder_context):
    return Fruit()
```
instructs tf.function to trace with the `Fruit()` objects
instead of the actual `Apple()` and `Mango()` objects when it receives a
call to `get_mixed_flavor(Apple(), Mango())`. For example, Tensor arguments
are replaced with Tensors of similar shape and dtype, output from
a tf.Placeholder op.

More generally, placeholder values are the arguments of a tf.function,
as seen from the function's body:
```python
@tf.function
def foo(x):
  # Here `x` is be the placeholder value
  ...

foo(x) # Here `x` is the actual value
```


def _get_model_info(func, parent_class):
    from transformers.models import auto as auto_module
    if parent_class is not None:
        model_name_lowercase = get_model_name(parent_class)
    else:
        model_name_lowercase = get_model_name(func)
    if model_name_lowercase and model_name_lowercase not in getattr(getattr(auto_module, PLACEHOLDER_TO_AUTO_MODULE['config_class'][0]), PLACEHOLDER_TO_AUTO_MODULE['config_class'][1]):
        model_name_lowercase = model_name_lowercase.replace('_', '-')
    class_name = func.__qualname__.split('.')[0]
    if model_name_lowercase is None:
        config_class = None
    else:
        try:
            config_class = getattr(getattr(auto_module, PLACEHOLDER_TO_AUTO_MODULE['config_class'][0]), PLACEHOLDER_TO_AUTO_MODULE['config_class'][1])[model_name_lowercase]
        except KeyError:
            if model_name_lowercase in HARDCODED_CONFIG_FOR_MODELS:
                config_class = HARDCODED_CONFIG_FOR_MODELS[model_name_lowercase]
            else:
                config_class = 'ModelConfig'
                print(f'🚨 Config not found for {model_name_lowercase}. You can manually add it to HARDCODED_CONFIG_FOR_MODELS in utils/args_doc.py')
    return (model_name_lowercase, class_name, config_class)

Extract model information from a function or its parent class.

Args:
    func (`function`): The function to extract information from
    parent_class (`class`): Optional parent class of the function


class FlaxCLIPTextModelOutput(ModelOutput):
    text_embeds: jnp.ndarray = None
    last_hidden_state: jnp.ndarray = None
    hidden_states: Optional[Tuple[jnp.ndarray, ...]] = None
    attentions: Optional[Tuple[jnp.ndarray, ...]] = None

Base class for text model's outputs that also contains a pooling of the last hidden states.

Args:
    text_embeds (`jnp.ndarray` of shape `(batch_size, output_dim`):
        The text embeddings obtained by applying the projection layer to the pooled output of
        [`FlaxCLIPTextModel`].
    last_hidden_state (`jnp.ndarray` of shape `(batch_size, sequence_length, hidden_size)`):
        Sequence of hidden-states at the output of the last layer of the model.
    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 get_shape(x: tf.Tensor, name: Optional[str]=None) -> Union[tf.TensorShape, types.IntTensor]:
    name = 'get_shape' if name is None else name
    with tf.name_scope(name):
        x = tf.convert_to_tensor(x)
        is_fully_defined = x.shape.is_fully_defined()
        if is_fully_defined:
            return x.shape
        return tf.shape(x)

Returns static shape of `x` if it is fully defined, or dynamic, otherwise.

####Example
```python
import tensorflow as tf
import tf_quant_finance as tff

x = tf.zeros([5, 2])
prefer_static_shape(x)
# Expected: [5, 2]

Args:
  x: A tensor of any shape and `dtype`
  name: Python string. The name to give to the ops created by this function.
    Default value: `None` which maps to the default name
    `get_shape`.

Returns:
  A shape of `x` which a list, if the shape is fully defined, or a `Tensor`
  for dynamically shaped `x`.


class TensorBoardCallback(TrainerCallback):

    def __init__(self, tb_writer=None):
        has_tensorboard = is_tensorboard_available()
        if not has_tensorboard:
            raise RuntimeError('TensorBoardCallback requires tensorboard to be installed. Either update your PyTorch version or install tensorboardX.')
        if has_tensorboard:
            try:
                from torch.utils.tensorboard import SummaryWriter
                self._SummaryWriter = SummaryWriter
            except ImportError:
                try:
                    from tensorboardX import SummaryWriter
                    self._SummaryWriter = SummaryWriter
                except ImportError:
                    self._SummaryWriter = None
        else:
            self._SummaryWriter = None
        self.tb_writer = tb_writer

    def _init_summary_writer(self, args, log_dir=None):
        log_dir = log_dir or args.logging_dir
        if self._SummaryWriter is not None:
            self.tb_writer = self._SummaryWriter(log_dir=log_dir)

    def on_train_begin(self, args, state, control, **kwargs):
        if not state.is_world_process_zero:
            return
        log_dir = None
        if state.is_hyper_param_search:
            trial_name = state.trial_name
            if trial_name is not None:
                log_dir = os.path.join(args.logging_dir, trial_name)
        if self.tb_writer is None:
            self._init_summary_writer(args, log_dir)
        if self.tb_writer is not None:
            self.tb_writer.add_text('args', args.to_json_string())
            if 'model' in kwargs:
                model = kwargs['model']
                if hasattr(model, 'config') and model.config is not None:
                    model_config_json = model.config.to_json_string()
                    self.tb_writer.add_text('model_config', model_config_json)

    def on_log(self, args, state, control, logs=None, **kwargs):
        if not state.is_world_process_zero:
            return
        if self.tb_writer is None:
            self._init_summary_writer(args)
        if self.tb_writer is not None:
            logs = rewrite_logs(logs)
            for k, v in logs.items():
                if isinstance(v, (int, float)):
                    self.tb_writer.add_scalar(k, v, state.global_step)
                elif isinstance(v, str):
                    self.tb_writer.add_text(k, v, state.global_step)
                else:
                    logger.warning(f'''Trainer is attempting to log a value of "{v}" of type {type(v)} for key "{k}" as a scalar. This invocation of Tensorboard's writer.add_scalar() is incorrect so we dropped this attribute.''')
            self.tb_writer.flush()

    def on_train_end(self, args, state, control, **kwargs):
        if self.tb_writer:
            self.tb_writer.close()
            self.tb_writer = None

A [`TrainerCallback`] that sends the logs to [TensorBoard](https://www.tensorflow.org/tensorboard).

Args:
    tb_writer (`SummaryWriter`, *optional*):
        The writer to use. Will instantiate one if not set.


def from_config(cls, config):
    return cls(**config)

Instantiates an initializer from a configuration dictionary.

Example:

```python
initializer = RandomUniform(-1, 1)
config = initializer.get_config()
initializer = RandomUniform.from_config(config)
```

Args:
  config: A Python dictionary. It will typically be the output of
    `get_config`.

Returns:
  An Initializer instance.


def format_image_annotations_as_coco(image_id: str, categories: list[int], areas: list[float], bboxes: list[tuple[float]]) -> dict:
    annotations = []
    for category, area, bbox in zip(categories, areas, bboxes):
        formatted_annotation = {'image_id': image_id, 'category_id': category, 'iscrowd': 0, 'area': area, 'bbox': list(bbox)}
        annotations.append(formatted_annotation)
    return {'image_id': image_id, 'annotations': annotations}

Format one set of image annotations to the COCO format

Args:
    image_id (str): image id. e.g. "0001"
    categories (List[int]): list of categories/class labels corresponding to provided bounding boxes
    areas (List[float]): list of corresponding areas to provided bounding boxes
    bboxes (List[Tuple[float]]): list of bounding boxes provided in COCO format
        ([center_x, center_y, width, height] in absolute coordinates)

Returns:
    dict: {
        "image_id": image id,
        "annotations": list of formatted annotations
    }


def real(input, name=None):
    with ops.name_scope(name, 'Real', [input]) as name:
        input = ops.convert_to_tensor(input, name='input')
        if input.dtype.is_complex:
            real_dtype = input.dtype.real_dtype
            return gen_math_ops.real(input, Tout=real_dtype, name=name)
        elif input.dtype.is_numeric:
            return input
        else:
            raise TypeError('input must be a numeric tensor, but got tensor with dtype {}'.format(input.dtype))

Returns the real part of a complex (or real) tensor.

Given a tensor `input`, this operation returns a tensor of type `float` that
is the real part of each element in `input` considered as a complex number.

For example:

```python
x = tf.constant([-2.25 + 4.75j, 3.25 + 5.75j])
tf.math.real(x)  # [-2.25, 3.25]
```

If `input` is already real, it is returned unchanged.

Args:
  input: A `Tensor`. Must have numeric type.
  name: A name for the operation (optional).

Returns:
  A `Tensor` of type `float32` or `float64`.


def tokenize(self, text, never_split=None):
    never_split = self.never_split.union(set(never_split)) if never_split else self.never_split
    text = self._clean_text(text)
    if self.tokenize_chinese_chars:
        text = self._tokenize_chinese_chars(text)
    unicode_normalized_text = unicodedata.normalize('NFC', text)
    orig_tokens = whitespace_tokenize(unicode_normalized_text)
    split_tokens = []
    for token in orig_tokens:
        if token not in never_split:
            if self.do_lower_case:
                token = token.lower()
                if self.strip_accents is not False:
                    token = self._run_strip_accents(token)
            elif self.strip_accents:
                token = self._run_strip_accents(token)
        split_tokens.extend(self._run_split_on_punc(token, never_split))
    output_tokens = whitespace_tokenize(' '.join(split_tokens))
    return output_tokens

Basic Tokenization of a piece of text. For sub-word tokenization, see WordPieceTokenizer.

Args:
    never_split (`List[str]`, *optional*)
        Kept for backward compatibility purposes. Now implemented directly at the base class level (see
        [`PreTrainedTokenizer.tokenize`]) List of token not to split.


def _GetActualMessage(self):
    if six.PY2:
        return self._actual.message
    return self._actual.args[0] if self._actual.args else ''

Returns the "message" portion of an exception.

Many Python 2 exceptions have a "message" attribute, so return that directly
in Python 2. However, this attribute is never present in Python 3, so return
the first argument passed to the exception instance as the message.

Returns:
  String


def scatter_nd_update(ref, indices, updates, use_locking=True, name=None):
    if ref.dtype._is_ref_dtype:
        return gen_state_ops.scatter_nd_update(ref, indices, updates, use_locking, name)
    return ref._lazy_read(gen_state_ops.resource_scatter_nd_update(ref.handle, indices, ops.convert_to_tensor(updates, ref.dtype), name=name))

Applies sparse `updates` to individual values or slices in a Variable.

`ref` is a `Tensor` with rank `P` and `indices` is a `Tensor` of rank `Q`.

`indices` must be integer tensor, containing indices into `ref`.
It must be shape `[d_0, ..., d_{Q-2}, K]` where `0 < K <= P`.

The innermost dimension of `indices` (with length `K`) corresponds to
indices into elements (if `K = P`) or slices (if `K < P`) along the `K`th
dimension of `ref`.

`updates` is `Tensor` of rank `Q-1+P-K` with shape:

```
[d_0, ..., d_{Q-2}, ref.shape[K], ..., ref.shape[P-1]].
```

For example, say we want to update 4 scattered elements to a rank-1 tensor to
8 elements. In Python, that update would look like this:

```python
    ref = tf.Variable([1, 2, 3, 4, 5, 6, 7, 8])
    indices = tf.constant([[4], [3], [1] ,[7]])
    updates = tf.constant([9, 10, 11, 12])
    update = tf.compat.v1.scatter_nd_update(ref, indices, updates)
    with tf.compat.v1.Session() as sess:
      print sess.run(update)
```

The resulting update to ref would look like this:

    [1, 11, 3, 10, 9, 6, 7, 12]

See `tf.scatter_nd` for more details about how to make updates to
slices.

Args:
  ref: A Variable.
  indices: A `Tensor`. Must be one of the following types: `int32`, `int64`.
    A tensor of indices into ref.
  updates: A `Tensor`. Must have the same type as `ref`.
    A Tensor. Must have the same type as ref. A tensor of updated
    values to add to ref.
  use_locking: An optional `bool`. Defaults to `True`.
    An optional bool. Defaults to True. If True, the assignment will
    be protected by a lock; otherwise the behavior is undefined,
    but may exhibit less contention.
  name: A name for the operation (optional).

Returns:
  The value of the variable after the update.


def register(self, candidate, name=None):
    if not name:
        name = candidate.__name__
    if name in self._registry:
        frame = self._registry[name][_LOCATION_TAG]
        raise KeyError("Registering two %s with name '%s'! (Previous registration was in %s %s:%d)" % (self._name, name, frame.name, frame.filename, frame.lineno))
    logging.vlog(1, 'Registering %s (%s) in %s.', name, candidate, self._name)
    stack = traceback.extract_stack(limit=3)
    stack_index = min(2, len(stack) - 1)
    if stack_index >= 0:
        location_tag = stack[stack_index]
    else:
        location_tag = ('UNKNOWN', 'UNKNOWN', 'UNKNOWN', 'UNKNOWN', 'UNKNOWN')
    self._registry[name] = {_TYPE_TAG: candidate, _LOCATION_TAG: location_tag}

Registers a Python object "candidate" for the given "name".

Args:
  candidate: The candidate object to add to the registry.
  name: An optional string specifying the registry key for the candidate.
        If None, candidate.__name__ will be used.
Raises:
  KeyError: If same name is used twice.


def get_reachable_from_inputs(inputs, targets=None):
    inputs = nest.flatten(inputs, expand_composites=True)
    reachable = object_identity.ObjectIdentitySet(inputs)
    if targets:
        remaining_targets = object_identity.ObjectIdentitySet(nest.flatten(targets))
    queue = collections.deque(inputs)
    while queue:
        x = queue.pop()
        if isinstance(x, tuple(_user_convertible_tensor_types)):
            continue
        if isinstance(x, ops.Operation):
            outputs = x.outputs[:] or []
            outputs += x._control_outputs
        elif isinstance(x, variables.Variable):
            try:
                outputs = [x.op]
            except AttributeError:
                outputs = []
        elif tensor_util.is_tf_type(x):
            outputs = x.consumers()
        else:
            raise TypeError('Expected Operation, Variable, or Tensor, got ' + str(x))
        for y in outputs:
            if y not in reachable:
                reachable.add(y)
                if targets:
                    remaining_targets.discard(y)
                queue.appendleft(y)
        if targets and (not remaining_targets):
            return reachable
    return reachable

Returns the set of tensors/ops reachable from `inputs`.

Stops if all targets have been found (target is optional).

Only valid in Symbolic mode, not Eager mode.

Args:
  inputs: List of tensors.
  targets: List of tensors.

Returns:
  A set of tensors reachable from the inputs (includes the inputs themselves).


def __new__(cls, *args) -> 'InvokeExpressionNode':
    if not args:
        return super().__new__(cls)
    _, identifier, parent_node = args
    if identifier == 'reference' and isinstance(parent_node.return_type, _fhir_path_data_types.ReferenceStructureDataType):
        return super().__new__(InvokeReferenceNode)
    return super().__new__(InvokeExpressionNode)

Creates a new InvokeExpressionNode node or one of its subclasses.

Creates either an InvokeExpressionNode or InvokeReferenceNode, a subclass of
InvokeExpressionNode. The InvokeReferenceNode is returned when a field named
'reference' is invoked against a FHIR Reference resource. Database backends
have special behavior for reference nodes. This reference-specific node type
allows them to define visitors to implement their reference-specific logic.

Args:
  *args: The args passed to `__init__`.

Returns:
  A new InvokeExpressionNode of the appropriate type.


def is_slice_on_extension(element_definition: message.Message) -> bool:
    type_codes = element_type_codes(element_definition)
    return 'Extension' in type_codes and _SLICE_ON_EXTENSION_ID_RE.search(cast(Any, element_definition).id.value) is not None

Returns `True` if the given element is describing a slice on an extension.

More information about extensions:
http://hl7.org/fhir/defining-extensions.html.

Args:
  element_definition: The element definition (element) that we are checking.


def floor_divide(x1, x2):
    if any_symbolic_tensors((x1, x2)):
        return FloorDivide().symbolic_call(x1, x2)
    return backend.numpy.floor_divide(x1, x2)

Returns the largest integer smaller or equal to the division of inputs.

Args:
    x1: Numerator.
    x2: Denominator.

Returns:
    Output tensor, `y = floor(x1/x2)`


def forward(self, input_ids: Optional[torch.Tensor]=None, attention_mask: Optional[torch.Tensor]=None, position_ids: Optional[torch.Tensor]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None) -> Union[Tuple, Blip2TextModelOutput]:
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    query_embeds = self.embeddings(input_ids=input_ids, position_ids=position_ids)
    text_outputs = self.qformer(query_embeds=query_embeds, query_length=0, attention_mask=attention_mask, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
    pooled_output = text_outputs[0] if not return_dict else text_outputs.last_hidden_state
    pooled_output = pooled_output.to(dtype=self.text_projection.weight.dtype)
    text_embeds = self.text_projection(pooled_output)
    text_embeds = nn.functional.normalize(text_embeds, dim=-1)
    if not return_dict:
        outputs = (text_embeds, text_outputs[0]) + text_outputs[2:]
        return tuple((output for output in outputs if output is not None))
    return Blip2TextModelOutput(text_embeds=text_embeds, last_hidden_state=text_outputs.last_hidden_state, hidden_states=text_outputs.hidden_states, attentions=text_outputs.attentions)

Examples:

```python
>>> import torch
>>> from transformers import AutoProcessor, Blip2TextModelWithProjection

>>> device = "cuda" if torch.cuda.is_available() else "cpu"

>>> model = Blip2TextModelWithProjection.from_pretrained(
...     "Salesforce/blip2-itm-vit-g", torch_dtype=torch.float16
... )

>>> model.to(device)  # doctest: +IGNORE_RESULT

>>> processor = AutoProcessor.from_pretrained("Salesforce/blip2-itm-vit-g")

>>> inputs = processor(text=["a photo of a cat", "a photo of a dog"], return_tensors="pt").to(device)

>>> outputs = model(**inputs)
>>> text_embeds = outputs.text_embeds
>>> print(text_embeds.shape)
torch.Size([2, 7, 256])
```


Called at the end of training.

Subclasses should override for any actions to run.

Args:
    logs: Dict. Currently the output of the last call to `on_epoch_end()`
      is passed to this argument for this method but that may change in
      the future.


def _clone_functional_model(model, input_tensors=None, layer_fn=_clone_layer):
    if not isinstance(model, Model):
        raise ValueError('Expected `model` argument to be a `Model` instance, got ', model)
    if isinstance(model, Sequential):
        raise ValueError('Expected `model` argument to be a functional `Model` instance, got a `Sequential` instance instead:', model)
    if not model._is_graph_network:
        raise ValueError('Expected `model` argument to be a functional `Model` instance, but got a subclass model instead.')
    new_input_layers = {}
    if input_tensors is not None:
        input_tensors = nest.flatten(input_tensors)
        for i, input_tensor in enumerate(input_tensors):
            original_input_layer = model._input_layers[i]
            if not backend.is_keras_tensor(input_tensor):
                name = original_input_layer.name
                input_tensor = Input(tensor=input_tensor, name='input_wrapper_for_' + name)
                newly_created_input_layer = input_tensor._keras_history.layer
                new_input_layers[original_input_layer] = newly_created_input_layer
            else:
                new_input_layers[original_input_layer] = original_input_layer
    if not callable(layer_fn):
        raise ValueError('Expected `layer_fn` argument to be a callable.')
    model_configs, created_layers = _clone_layers_and_model_config(model, new_input_layers, layer_fn)
    input_tensors, output_tensors, created_layers = functional.reconstruct_from_config(model_configs, created_layers=created_layers)
    metrics_names = model.metrics_names
    model = Model(input_tensors, output_tensors, name=model.name)
    ancillary_layers = [layer for layer in created_layers.values() if layer not in model.layers]
    if ancillary_layers:
        new_nodes = nest.flatten([layer.inbound_nodes[1:] if functional._should_skip_first_node(layer) else layer.inbound_nodes for layer in created_layers.values()])
        _insert_ancillary_layers(model, ancillary_layers, metrics_names, new_nodes)
    return model

Clone a functional `Model` instance.

Model cloning is similar to calling a model on new inputs,
except that it creates new layers (and thus new weights) instead
of sharing the weights of the existing layers.

Input layers are always cloned.

Args:
    model: Instance of `Model`.
    input_tensors: optional list of input tensors
        to build the model upon. If not provided,
        placeholders will be created.
    layer_fn: callable to be applied on non-input layers in the model. By
        default it clones the layer. Another example is to preserve the layer
        to share the weights. This is required when we create a per-replica
        copy of the model with distribution strategy; we want the weights to
        be shared but still feed inputs separately so we create new input
        layers.

Returns:
    An instance of `Model` reproducing the behavior
    of the original model, on top of new inputs tensors,
    using newly instantiated weights.

Raises:
    ValueError: in case of invalid `model` argument value or `layer_fn`
    argument value.


def imag(input, name=None):
    with ops.name_scope(name, 'Imag', [input]) as name:
        input = ops.convert_to_tensor(input, name='input')
        if input.dtype.is_complex:
            return gen_math_ops.imag(input, Tout=input.dtype.real_dtype, name=name)
        else:
            return array_ops.zeros_like(input)

Returns the imaginary part of a complex (or real) tensor.

Given a tensor `input`, this operation returns a tensor of type `float` that
is the imaginary part of each element in `input` considered as a complex
number. If `input` is real, a tensor of all zeros is returned.

For example:

```python
x = tf.constant([-2.25 + 4.75j, 3.25 + 5.75j])
tf.math.imag(x)  # [4.75, 5.75]
```

Args:
  input: A `Tensor`. Must be one of the following types: `float`, `double`,
    `complex64`, `complex128`.
  name: A name for the operation (optional).

Returns:
  A `Tensor` of type `float32` or `float64`.


def forward(self, pixel_values: torch.FloatTensor, pixel_mask: Optional[torch.LongTensor]=None, encoder_outputs: Optional[torch.FloatTensor]=None, inputs_embeds: Optional[torch.FloatTensor]=None, decoder_inputs_embeds: Optional[torch.FloatTensor]=None, labels: Optional[List[dict]]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, **loss_kwargs) -> Union[Tuple[torch.FloatTensor], RTDetrV2ObjectDetectionOutput]:
    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
    output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    outputs = self.model(pixel_values, pixel_mask=pixel_mask, encoder_outputs=encoder_outputs, inputs_embeds=inputs_embeds, decoder_inputs_embeds=decoder_inputs_embeds, labels=labels, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
    denoising_meta_values = outputs.denoising_meta_values if return_dict else outputs[-1] if self.training else None
    outputs_class = outputs.intermediate_logits if return_dict else outputs[2]
    outputs_coord = outputs.intermediate_reference_points if return_dict else outputs[3]
    predicted_corners = outputs.intermediate_predicted_corners if return_dict else outputs[4]
    initial_reference_points = outputs.initial_reference_points if return_dict else outputs[5]
    logits = outputs_class[:, -1]
    pred_boxes = outputs_coord[:, -1]
    loss, loss_dict, auxiliary_outputs, enc_topk_logits, enc_topk_bboxes = (None, None, None, None, None)
    if labels is not None:
        enc_topk_logits = outputs.enc_topk_logits if return_dict else outputs[-5]
        enc_topk_bboxes = outputs.enc_topk_bboxes if return_dict else outputs[-4]
        loss, loss_dict, auxiliary_outputs = self.loss_function(logits, labels, self.device, pred_boxes, self.config, outputs_class, outputs_coord, enc_topk_logits=enc_topk_logits, enc_topk_bboxes=enc_topk_bboxes, denoising_meta_values=denoising_meta_values, predicted_corners=predicted_corners, initial_reference_points=initial_reference_points, **loss_kwargs)
    if not return_dict:
        if auxiliary_outputs is not None:
            output = (logits, pred_boxes) + (auxiliary_outputs,) + outputs
        else:
            output = (logits, pred_boxes) + outputs
        return (loss, loss_dict) + output if loss is not None else output
    return RTDetrV2ObjectDetectionOutput(loss=loss, loss_dict=loss_dict, logits=logits, pred_boxes=pred_boxes, auxiliary_outputs=auxiliary_outputs, last_hidden_state=outputs.last_hidden_state, intermediate_hidden_states=outputs.intermediate_hidden_states, intermediate_logits=outputs.intermediate_logits, intermediate_reference_points=outputs.intermediate_reference_points, intermediate_predicted_corners=outputs.intermediate_predicted_corners, initial_reference_points=outputs.initial_reference_points, decoder_hidden_states=outputs.decoder_hidden_states, decoder_attentions=outputs.decoder_attentions, cross_attentions=outputs.cross_attentions, encoder_last_hidden_state=outputs.encoder_last_hidden_state, encoder_hidden_states=outputs.encoder_hidden_states, encoder_attentions=outputs.encoder_attentions, init_reference_points=outputs.init_reference_points, enc_topk_logits=outputs.enc_topk_logits, enc_topk_bboxes=outputs.enc_topk_bboxes, enc_outputs_class=outputs.enc_outputs_class, enc_outputs_coord_logits=outputs.enc_outputs_coord_logits, denoising_meta_values=outputs.denoising_meta_values)

inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
    Optionally, instead of passing the flattened feature map (output of the backbone + projection layer), you
    can choose to directly pass a flattened representation of an image.
decoder_inputs_embeds (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_size)`, *optional*):
    Optionally, instead of initializing the queries with a tensor of zeros, you can choose to directly pass an
    embedded representation.
labels (`List[Dict]` of len `(batch_size,)`, *optional*):
    Labels for computing the bipartite matching loss. List of dicts, each dictionary containing at least the
    following 2 keys: 'class_labels' and 'boxes' (the class labels and bounding boxes of an image in the batch
    respectively). The class labels themselves should be a `torch.LongTensor` of len `(number of bounding boxes
    in the image,)` and the boxes a `torch.FloatTensor` of shape `(number of bounding boxes in the image, 4)`.

Examples:

```python
>>> from transformers import RTDetrV2ImageProcessor, RTDetrV2ForObjectDetection
>>> from PIL import Image
>>> import requests
>>> import torch

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

>>> image_processor = RTDetrV2ImageProcessor.from_pretrained("PekingU/RTDetrV2_r50vd")
>>> model = RTDetrV2ForObjectDetection.from_pretrained("PekingU/RTDetrV2_r50vd")

>>> # prepare image for the model
>>> inputs = image_processor(images=image, return_tensors="pt")

>>> # forward pass
>>> outputs = model(**inputs)

>>> logits = outputs.logits
>>> list(logits.shape)
[1, 300, 80]

>>> boxes = outputs.pred_boxes
>>> list(boxes.shape)
[1, 300, 4]

>>> # convert outputs (bounding boxes and class logits) to Pascal VOC format (xmin, ymin, xmax, ymax)
>>> target_sizes = torch.tensor([image.size[::-1]])
>>> results = image_processor.post_process_object_detection(outputs, threshold=0.9, target_sizes=target_sizes)[
...     0
... ]

>>> for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
...     box = [round(i, 2) for i in box.tolist()]
...     print(
...         f"Detected {model.config.id2label[label.item()]} with confidence "
...         f"{round(score.item(), 3)} at location {box}"
...     )
Detected sofa with confidence 0.97 at location [0.14, 0.38, 640.13, 476.21]
Detected cat with confidence 0.96 at location [343.38, 24.28, 640.14, 371.5]
Detected cat with confidence 0.958 at location [13.23, 54.18, 318.98, 472.22]
Detected remote with confidence 0.951 at location [40.11, 73.44, 175.96, 118.48]
Detected remote with confidence 0.924 at location [333.73, 76.58, 369.97, 186.99]
```


def deserialize(self, stamp_token, serialized_proto):
    return gen_boosted_trees_ops.boosted_trees_deserialize_ensemble(self.resource_handle, stamp_token, serialized_proto)

Deserialize the input proto and resets the ensemble from it.

Args:
  stamp_token: int64 scalar Tensor to denote the stamp of the resource.
  serialized_proto: string scalar Tensor of the serialized proto.

Returns:
  Operation (for dependencies).


def _current_graph(op_input_list, graph=None):
    current_default_graph = ops.get_default_graph()
    if current_default_graph.building_function:
        return current_default_graph
    op_input_list = tuple(op_input_list)
    if graph and (not isinstance(graph, ops.Graph)):
        raise TypeError('Input graph needs to be a Graph: %s' % (graph,))
    original_graph_element = None
    for op_input in op_input_list:
        if isinstance(op_input, (ops.Operation, tensor_lib.Tensor, composite_tensor.CompositeTensor)) and (not isinstance(op_input, tensor_lib.Tensor) or type(op_input) == tensor_lib.Tensor):
            graph_element = op_input
        else:
            graph_element = _as_graph_element(op_input)
        if graph_element is not None:
            if not graph:
                original_graph_element = graph_element
                graph = getattr(graph_element, 'graph', None)
            elif original_graph_element is not None:
                _assert_same_graph(original_graph_element, graph_element)
            elif graph_element.graph is not graph:
                raise ValueError('%s is not from the passed-in graph.' % graph_element)
    return graph or current_default_graph

Returns the appropriate graph to use for the given inputs.

This library method provides a consistent algorithm for choosing the graph
in which an Operation should be constructed:

1. If the default graph is being used to construct a function, we
   use the default graph.
2. If the "graph" is specified explicitly, we validate that all of the inputs
   in "op_input_list" are compatible with that graph.
3. Otherwise, we attempt to select a graph from the first Operation-
   or Tensor-valued input in "op_input_list", and validate that all other
   such inputs are in the same graph.
4. If the graph was not specified and it could not be inferred from
   "op_input_list", we attempt to use the default graph.

Args:
  op_input_list: A list of inputs to an operation, which may include `Tensor`,
    `Operation`, and other objects that may be converted to a graph element.
  graph: (Optional) The explicit graph to use.

Raises:
  TypeError: If op_input_list is not a list or tuple, or if graph is not a
    Graph.
  ValueError: If a graph is explicitly passed and not all inputs are from it,
    or if the inputs are from multiple graphs, or we could not find a graph
    and there was no default graph.

Returns:
  The appropriate graph to use for the given inputs.


def plot_model(model, to_file='model.png', show_shapes=False, show_dtype=False, show_layer_names=False, rankdir='TB', expand_nested=False, dpi=200, show_layer_activations=False, show_trainable=False, **kwargs):
    if not model.built:
        raise ValueError('This model has not yet been built. Build the model first by calling `build()` or by calling the model on a batch of data.')
    if not check_pydot():
        message = 'You must install pydot (`pip install pydot`) for `plot_model` to work.'
        if 'IPython.core.magics.namespace' in sys.modules:
            io_utils.print_msg(message)
            return
        else:
            raise ImportError(message)
    if not check_graphviz():
        message = 'You must install graphviz (see instructions at https://graphviz.gitlab.io/download/) for `plot_model` to work.'
        if 'IPython.core.magics.namespace' in sys.modules:
            io_utils.print_msg(message)
            return
        else:
            raise ImportError(message)
    if kwargs.pop('layer_range', None) is not None:
        raise ValueError('Argument `layer_range` is no longer supported.')
    if kwargs:
        raise ValueError(f'Unrecognized keyword arguments: {kwargs}')
    dot = model_to_dot(model, show_shapes=show_shapes, show_dtype=show_dtype, show_layer_names=show_layer_names, rankdir=rankdir, expand_nested=expand_nested, dpi=dpi, show_layer_activations=show_layer_activations, show_trainable=show_trainable)
    to_file = str(to_file)
    if dot is None:
        return
    _, extension = os.path.splitext(to_file)
    if not extension:
        extension = 'png'
    else:
        extension = extension[1:]
    dot.write(to_file, format=extension)
    if extension != 'pdf':
        try:
            from IPython import display
            return display.Image(filename=to_file)
        except ImportError:
            pass

Converts a Keras model to dot format and save to a file.

Example:

```python
inputs = ...
outputs = ...
model = keras.Model(inputs=inputs, outputs=outputs)

dot_img_file = '/tmp/model_1.png'
keras.utils.plot_model(model, to_file=dot_img_file, show_shapes=True)
```

Args:
    model: A Keras model instance
    to_file: File name of the plot image.
    show_shapes: whether to display shape information.
    show_dtype: whether to display layer dtypes.
    show_layer_names: whether to display layer names.
    rankdir: `rankdir` argument passed to PyDot,
        a string specifying the format of the plot: `"TB"`
        creates a vertical plot; `"LR"` creates a horizontal plot.
    expand_nested: whether to expand nested Functional models
        into clusters.
    dpi: Image resolution in dots per inch.
    show_layer_activations: Display layer activations (only for layers that
        have an `activation` property).
    show_trainable: whether to display if a layer is trainable.

Returns:
    A Jupyter notebook Image object if Jupyter is installed.
    This enables in-line display of the model plots in notebooks.


class AriaSharedExpertsMLP(nn.Module):

    def __init__(self, config: AriaTextConfig):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size * config.moe_num_shared_experts
        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias)
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, x):
        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
        return down_proj

Shared Expert MLP for shared experts.

Unlike routed experts, shared experts process all tokens without routing.
This class reconfigures the intermediate size in comparison to the LlamaMLP.

Args:
    config (`AriaTextConfig`): Configuration object for the Aria language model.


def _check_params(window_length, dtype):
    if not dtype.is_floating:
        raise ValueError('dtype must be a floating point type. Found %s' % dtype)
    window_length = ops.convert_to_tensor(window_length, dtype=dtypes.int32)
    window_length.shape.assert_has_rank(0)
    return window_length

Check window_length and dtype params.

Args:
  window_length: A scalar value or `Tensor`.
  dtype: The data type to produce. Must be a floating point type.

Returns:
  window_length converted to a tensor of type int32.

Raises:
  ValueError: If `dtype` is not a floating point type or window_length is not
    a scalar.


def _prefix_output_keys(self, output_dict, output_name):
    new_outputs = {}
    for key, val in output_dict.items():
        key = self._prefix_key(key, output_name)
        new_outputs[key] = val
    return new_outputs

Prepend output_name to the output_dict keys if it doesn't exist.

This produces predictable prefixes for the pre-determined outputs
of SupervisedOutput.

Args:
  output_dict: dict of string to Tensor, assumed valid.
  output_name: prefix string to prepend to existing keys.

Returns:
  dict with updated keys and existing values.


def is_chief(cluster_spec=None, task_type=None, task_id=None):
    if has_worker_context():
        return dc_context.get_current_worker_context().is_chief
    _validate_cluster_spec(cluster_spec, task_type, task_id)
    cluster_spec = normalize_cluster_spec(cluster_spec).as_dict()
    if task_type == 'chief' or task_type == 'evaluator':
        return True
    if 'chief' not in cluster_spec and task_type == 'worker' and (task_id == 0):
        return True
    return False

Returns whether the given task is chief in the cluster.

Since there is at most one evaluator and the evaluator itself should be
independent of the training cluster, the evaluator job is also a chief job on
its own.

If this is currently running under a `_WorkerContext` of distribute
coordinator, the arguments can be omitted as the result is already available.

Args:
  cluster_spec: a dict, `ClusterDef` or `ClusterSpec` object specifying the
    cluster configurations.
  task_type: the task type in the cluster.
  task_id: the task id in the cluster.

Returns:
  a boolean indicating whether the given task is chief.

Raises:
  ValueError: if `task_type` is not in the `cluster_spec` or `task_id` exceeds
    the maximum id of the `task_type`.


def tokenize(self, text, never_split=None):
    never_split = self.never_split.union(set(never_split)) if never_split else self.never_split
    text = self._clean_text(text)
    if self.tokenize_chinese_chars:
        text = self._tokenize_chinese_chars(text)
    unicode_normalized_text = unicodedata.normalize('NFC', text)
    orig_tokens = whitespace_tokenize(unicode_normalized_text)
    split_tokens = []
    for token in orig_tokens:
        if token not in never_split:
            if self.do_lower_case:
                token = token.lower()
                if self.strip_accents is not False:
                    token = self._run_strip_accents(token)
            elif self.strip_accents:
                token = self._run_strip_accents(token)
        split_tokens.extend(self._run_split_on_punc(token, never_split))
    output_tokens = whitespace_tokenize(' '.join(split_tokens))
    return output_tokens

Basic Tokenization of a piece of text. For sub-word tokenization, see WordPieceTokenizer.

Args:
    never_split (`List[str]`, *optional*)
        Kept for backward compatibility purposes. Now implemented directly at the base class level (see
        [`PreTrainedTokenizer.tokenize`]) List of token not to split.


def _validate_fhir_constraints(msg: message.Message, base_name: str, primitive_handler_: primitive_handler.PrimitiveHandler) -> None:
    if annotation_utils.is_primitive_type(msg):
        _ = primitive_handler_.primitive_wrapper_from_primitive(msg)
        return
    if proto_utils.is_message_type(msg, any_pb2.Any):
        return
    for field in msg.DESCRIPTOR.fields:
        field_name = f'{base_name}.{proto_utils.json_field_name(field)}'
        _validate_field(msg, field, field_name, primitive_handler_)
    for oneof in msg.DESCRIPTOR.oneofs:
        if msg.WhichOneof(oneof.name) is None and (not oneof.GetOptions().HasExtension(annotations_pb2.fhir_oneof_is_optional)):
            raise fhir_errors.InvalidFhirError(f'Empty oneof: `{oneof.full_name}`.')

Iterates over fields of the provided message and validates constraints.

Args:
  msg: The message to validate.
  base_name: The root message name for recursive validation of nested message
    fields.
  primitive_handler_: Responsible for returning PrimitiveWrappers.

Raises:
  fhir_errors.InvalidFhirError: In the event that a field is found to be
  violating FHIR constraints or a required oneof is not set.


def _maybe_partial_apply_variables(fn, args, kwargs):

    def is_distributed_var(x):
        flat = nest.flatten(x)
        return flat and isinstance(flat[0], values.DistributedVariable)
    var_kwargs = {}
    nonvar_kwargs = {}
    if kwargs:
        var_kwargs = {k: v for k, v in kwargs.items() if is_distributed_var(v)}
    if var_kwargs:
        nonvar_kwargs = {k: v for k, v in kwargs.items() if not is_distributed_var(v)}
    positional_args = []
    index_of_star_args = None
    for i, p in enumerate(tf_inspect.signature(fn).parameters.values()):
        if i == 0 and p.name == 'self':
            continue
        if p.kind == tf_inspect.Parameter.POSITIONAL_OR_KEYWORD:
            positional_args.append(p.name)
        elif p.kind == tf_inspect.Parameter.VAR_POSITIONAL:
            index_of_star_args = i
        elif p.kind == tf_inspect.Parameter.POSITIONAL_ONLY:
            if var_kwargs or any((is_distributed_var(a) for a in args)):
                raise ValueError(f'Mixing Variables and positional-only parameters not supported by TPUStrategy. Received {len(var_kwargs)} DistributedVariables in **kwargs and {sum((is_distributed_var(a) for a in args))} in *args, expected zero for both.')
            return (fn, args, kwargs)
    star_args = []
    have_seen_var_arg = False
    for i, a in enumerate(args):
        if is_distributed_var(a):
            if index_of_star_args is not None and i >= index_of_star_args:
                raise ValueError('TPUStrategy.run() cannot handle Variables passed to *args. Either name the function argument, or capture the Variable implicitly.')
            if len(positional_args) <= i:
                raise ValueError('Too many positional arguments passed to call to TPUStrategy.run().')
            var_kwargs[positional_args[i]] = a
            have_seen_var_arg = True
        else:
            if index_of_star_args is not None and i >= index_of_star_args:
                if have_seen_var_arg:
                    raise ValueError('TPUStrategy.run() cannot handle both Variables and a mix of positional args and *args. Either remove the *args, or capture the Variable implicitly.')
                else:
                    star_args.append(a)
                    continue
            if len(positional_args) <= i:
                raise ValueError('Too many positional arguments passed to call to TPUStrategy.run().')
            nonvar_kwargs[positional_args[i]] = a
    if var_kwargs:
        return (functools.partial(fn, **var_kwargs), star_args, nonvar_kwargs)
    return (fn, args, kwargs)

Inspects arguments to partially apply any DistributedVariable.

This avoids an automatic cast of the current variable value to tensor.

Note that a variable may be captured implicitly with Python scope instead of
passing it to run(), but supporting run() keeps behavior consistent
with MirroredStrategy.

Since positional arguments must be applied from left to right, this function
does some tricky function inspection to move variable positional arguments
into kwargs. As a result of this, we can't support passing Variables as *args,
nor as args to functions which combine both explicit positional arguments and
*args.

Args:
  fn: The function to run, as passed to run().
  args: Positional arguments to fn, as passed to run().
  kwargs: Keyword arguments to fn, as passed to run().

Returns:
  A tuple of the function (possibly wrapped), args, kwargs (both
  possibly filtered, with members of args possibly moved to kwargs).
  If no variables are found, this function is a noop.

Raises:
  ValueError: If the function signature makes unsupported use of *args, or if
    too many arguments are passed.


def _prepare_grid(self, times, grid_step):
    grid = tf.range(0.0, times[-1], grid_step, dtype=self._dtype)
    all_times = tf.concat([grid, times], axis=0)
    mask = tf.concat([tf.zeros_like(grid, dtype=tf.bool), tf.ones_like(times, dtype=tf.bool)], axis=0)
    perm = tf.argsort(all_times, stable=True)
    all_times = tf.gather(all_times, perm)
    mask = tf.gather(mask, perm)
    return (all_times, mask)

Prepares grid of times for path generation.

Args:
  times:  Rank 1 `Tensor` of increasing positive real values. The times at
    which the path points are to be evaluated.
  grid_step: Rank 0 real `Tensor`. Maximal distance between points in
    resulting grid.

Returns:
  Tuple `(all_times, mask)`.
  `all_times` is 1-D real `Tensor` containing all points from 'times` and
  whose intervals are at most `grid_step`.
  `mask` is a boolean 1-D tensor of the same shape as 'all_times', showing
  which elements of 'all_times' correspond to values from `times`.
  Guarantees that times[0]=0 and grid_step[0]=False.
  'all_times` is sorted ascending and may contain duplicates.


def advise(graph=None, run_meta=None, options=_DEFAULT_ADVISE_OPTIONS):
    if not graph and (not context.executing_eagerly()):
        graph = ops.get_default_graph()
    if options == _DEFAULT_ADVISE_OPTIONS:
        options = ALL_ADVICE.copy()
    op_log = tfprof_logger.merge_default_with_oplog(graph, None, run_meta, add_trace=True)
    run_meta_str = run_meta.SerializeToString() if run_meta else b''
    opts = _build_advisor_options(options)
    ret = tfprof_output_pb2.AdviceProto()
    ret.ParseFromString(print_mdl.PrintModelAnalysis(_graph_string(graph), run_meta_str, op_log.SerializeToString(), 'advise'.encode('utf-8'), opts.SerializeToString()))
    return ret

Auto profile and advise.

  Builds profiles and automatically check anomalies of various
  aspects. For more details:
  https://github.com/tensorflow/tensorflow/tree/master/tensorflow/core/profiler/README.md

Args:
  graph: tf.Graph. If None and eager execution is not enabled, use default
    graph.
  run_meta: optional tensorflow.RunMetadata proto. It is necessary to
    support run time information profiling, such as time and memory.
  options: see ALL_ADVICE example above. Default checks everything.

Returns:
  Returns AdviceProto proto


def flatten(index, name='segmented_flatten'):
    batch_size = torch.prod(torch.tensor(list(index.batch_shape())))
    offset = torch.arange(start=0, end=batch_size, device=index.num_segments.device) * index.num_segments
    offset = offset.view(index.batch_shape())
    for _ in range(index.batch_dims, len(index.indices.size())):
        offset = offset.unsqueeze(-1)
    indices = offset + index.indices
    return IndexMap(indices=indices.view(-1), num_segments=index.num_segments * batch_size, batch_dims=0)

Flattens a batched index map (which is typically of shape batch_size, seq_length) to a 1d index map. This operation
relabels the segments to keep batch elements distinct. The k-th batch element will have indices shifted by
*num_segments* * (k - 1). The result is a tensor with *num_segments* multiplied by the number of elements in the
batch.

Args:
    index (`IndexMap`):
        IndexMap to flatten.
    name (`str`, *optional*, defaults to 'segmented_flatten'):
        Name for the operation. Currently not used

Returns:
    (`IndexMap`): The flattened IndexMap.


def build(self):

    def _create_per_worker_dataset():
        dataset = self._dataset_fn()
        return dataset
    per_worker_dataset = self._coordinator._create_per_worker_resources(_create_per_worker_dataset)
    dataset_fn_output_type_spec = self._dataset_fn.structured_outputs._type_spec
    for dataset_remote_value in per_worker_dataset._values:
        dataset_remote_value._type_spec = dataset_fn_output_type_spec
    return per_worker_dataset

Trigger dataset creation on workers without creating an iterator.

Returns:
  A PerWorkerValues object containing a tuple of RemoteValues, themselves
  containing the built Dataset for each worker


def infer_graph(inputs: Optional[Set[EventSetNode]], outputs: Set[EventSetNode]) -> Graph:
    graph = Graph()
    graph.outputs.update(outputs)
    pending_nodes: Set[EventSetNode] = outputs.copy()
    done_nodes: Set[EventSetNode] = set()
    missing_nodes: Set[EventSetNode] = set()
    while pending_nodes:
        node = next(iter(pending_nodes))
        pending_nodes.remove(node)
        assert node not in done_nodes
        graph.add_node(node)
        if inputs is not None and node in inputs:
            graph.inputs.add(node)
            continue
        if node.creator is None:
            if inputs is not None:
                missing_nodes.add(node)
            else:
                graph.inputs.add(node)
            continue
        graph.add_operator(node.creator)
        for input_node in node.creator.inputs.values():
            if input_node in done_nodes:
                continue
            pending_nodes.add(input_node)
        for output_node in node.creator.outputs.values():
            graph.add_node(output_node)
    if missing_nodes:
        raise ValueError(f'The following input nodes are required but not provided as input:\n{missing_nodes}')
    for e in graph.nodes:
        graph.add_sampling(e.sampling_node)
        for f in e.feature_nodes:
            graph.add_feature(f)
    return graph

Extracts the nodes in between the output and input nodes.

If inputs is set, fails if outputs cannot be computed from `inputs`.
If inputs is not set, infers the required set of inputs.

Args:
    inputs: Set of available input nodes. If None, inputs are inferred.
    outputs: Set of expected output nodes.

Returns:
    The inferred graph.

Raises:
    ValueError: If there are repeated nodes in the `inputs`; an
        unexpected type of input is provided; an unnamed node is inferred
        as input; or some nodes are required but not provided.


def sparse_intersection_indices_and_values(x1, x2):
    ones1 = tf.sparse.map_values(ones_like_int8, x1)
    ones2 = tf.sparse.map_values(ones_like_int8, x2)
    intersection_extra_dim = tf.sets.intersection(tf.sparse.expand_dims(ones1, axis=-1), tf.sparse.expand_dims(ones2, axis=-1))

    def empty_intersection():
        return (tf.zeros((0, x1.shape.rank), dtype=tf.int64), tf.zeros((0,), dtype=x1.values.dtype), tf.zeros((0,), dtype=x2.values.dtype))

    def non_empty_intersection():
        intersection = tf.sparse.reshape(intersection_extra_dim, x1.dense_shape)
        zeros1 = tf.sparse.map_values(zeros_like_int8, x1)
        zeros2 = tf.sparse.map_values(zeros_like_int8, x2)
        mask1 = tf.sparse.add(zeros1, intersection)
        mask2 = tf.sparse.add(zeros2, intersection)
        return (intersection.indices, tf.sparse.retain(x1, tf.cast(mask1.values, tf.bool)).values, tf.sparse.retain(x2, tf.cast(mask2.values, tf.bool)).values)
    return tf.cond(tf.equal(tf.size(intersection_extra_dim), 0), empty_intersection, non_empty_intersection)

Compute the indices for the intersection of two `tf.SparseTensor`s and
modify the values for these indices.

Args:
    x1: the first `tf.SparseTensor`.
    x2: the second `tf.SparseTensor`.
Returns: A tuple containing:
    - the indices for the intersection
    - `x1` values for the intersection indices (some values were removed)
    - `x2` values for the intersection indices (some values were removed)


def __init__(self, host: str, port: int, time_to_live: Union[int, timedelta], *, kwargs: Optional[Dict[str, Any]]=None, request_coder: Optional[coders.Coder], response_coder: Optional[coders.Coder], source_caller: Optional[Caller[RequestT, ResponseT]]=None):
    self.request_coder = request_coder
    self.response_coder = response_coder
    self.redis_caller = _RedisCaller(host, port, time_to_live, request_coder=self.request_coder, response_coder=self.response_coder, kwargs=kwargs, source_caller=source_caller, mode=_RedisMode.WRITE)

Args:
  host (str): The hostname or IP address of the Redis server.
  port (int): The port number of the Redis server.
  time_to_live: `(Union[int, timedelta])` The time-to-live (TTL) for
    records stored in Redis. Provide an integer (in seconds) or a
    `datetime.timedelta` object.
  kwargs: Optional(Dict[str, Any]) additional keyword arguments that
    are required to connect to your redis server. Same as `redis.Redis()`.
  request_coder: (Optional[`coders.Coder`]) coder for requests stored
    in Redis.
  response_coder: (Optional[`coders.Coder`]) coder for decoding responses
    received from Redis.
  source_caller: (Optional[`Caller`]): The source caller using this Redis
    cache in case of fetching the cache request to store in Redis.


def from_spec(cls, spec, name=None):
    return cls(spec.shape, spec.dtype, name or spec.name)

Returns a `TensorSpec` with the same shape and dtype as `spec`.

>>> spec = tf.TensorSpec(shape=[8, 3], dtype=tf.int32, name="OriginalName")
>>> tf.TensorSpec.from_spec(spec, "NewName")
TensorSpec(shape=(8, 3), dtype=tf.int32, name='NewName')

Args:
  spec: The `TypeSpec` used to create the new `TensorSpec`.
  name: The name for the new `TensorSpec`.  Defaults to `spec.name`.


class PerceiverProjectionPostprocessor(nn.Module):

    def __init__(self, in_channels: int, out_channels: int) -> None:
        super().__init__()
        self.classifier = nn.Linear(in_channels, out_channels)

    def forward(self, inputs: torch.Tensor, pos: Optional[torch.Tensor]=None, modality_sizes=None) -> torch.Tensor:
        logits = self.classifier(inputs)
        return logits

Projection postprocessing for Perceiver. Can be used to project the channels of the decoder output to a lower
dimension.

Args:
    in_channels (`int`):
        Number of channels in the input.
    out_channels (`int`):
        Number of channels in the output.


class QuantLinear(nn.Module):

    def __init__(self, in_features, out_features, bias=True, weight_bit=8, bias_bit=32, per_channel=False, quant_mode=False):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.weight = nn.Parameter(torch.zeros([out_features, in_features]))
        self.register_buffer('weight_integer', torch.zeros_like(self.weight))
        self.register_buffer('fc_scaling_factor', torch.zeros(self.out_features))
        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features))
            self.register_buffer('bias_integer', torch.zeros_like(self.bias))
        self.weight_bit = weight_bit
        self.quant_mode = quant_mode
        self.per_channel = per_channel
        self.bias_bit = bias_bit
        self.quant_mode = quant_mode
        self.percentile_mode = False
        self.weight_function = SymmetricQuantFunction.apply

    def __repr__(self):
        s = super().__repr__()
        s = f'({s} weight_bit={self.weight_bit}, quant_mode={self.quant_mode})'
        return s

    def forward(self, x, prev_act_scaling_factor=None):
        if not self.quant_mode:
            return (nn.functional.linear(x, weight=self.weight, bias=self.bias), None)
        assert prev_act_scaling_factor is not None and prev_act_scaling_factor.shape == (1,), 'Input activation to the QuantLinear layer should be globally (non-channel-wise) quantized. Please add a QuantAct layer with `per_channel = True` before this QuantAct layer'
        w = self.weight
        w_transform = w.data.detach()
        if self.per_channel:
            w_min, _ = torch.min(w_transform, dim=1, out=None)
            w_max, _ = torch.max(w_transform, dim=1, out=None)
        else:
            w_min = w_transform.min().expand(1)
            w_max = w_transform.max().expand(1)
        self.fc_scaling_factor = symmetric_linear_quantization_params(self.weight_bit, w_min, w_max, self.per_channel)
        self.weight_integer = self.weight_function(self.weight, self.weight_bit, self.percentile_mode, self.fc_scaling_factor)
        bias_scaling_factor = self.fc_scaling_factor * prev_act_scaling_factor
        if self.bias is not None:
            self.bias_integer = self.weight_function(self.bias, self.bias_bit, False, bias_scaling_factor)
        prev_act_scaling_factor = prev_act_scaling_factor.view(1, -1)
        x_int = x / prev_act_scaling_factor
        return (nn.functional.linear(x_int, weight=self.weight_integer, bias=self.bias_integer) * bias_scaling_factor, bias_scaling_factor)

Quantized version of `torch.nn.Linear`. Adds quantization-specific arguments on top of `torch.nn.Linear`.

Args:
    weight_bit (`int`, *optional*, defaults to `8`):
        Bitwidth for the quantized weight.
    bias_bit (`int`, *optional*, defaults to `32`):
        Bitwidth for the quantized bias.
    per_channel (`bool`, *optional*, defaults to `False`):
        Whether or not to use channel-wise quantization.
    quant_mode (`bool`, *optional*, defaults to `False`):
        Whether or not the layer is quantized.


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

List the jobs of a project. To list the jobs of a project in a region, we recommend using `projects.locations.jobs.list` with a [regional endpoint] (https://cloud.google.com/dataflow/docs/concepts/regional-endpoints). To list the all jobs across all regions, use `projects.jobs.aggregated`. Using `projects.jobs.list` is not recommended, because you can only get the list of jobs that are running in `us-central1`. `projects.locations.jobs.list` and `projects.jobs.list` support filtering the list of jobs by name. Filtering by name isn't supported by `projects.jobs.aggregated`.

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


def range_index_map(batch_shape, num_segments, name='range_index_map'):
    device = num_segments.device if torch.is_tensor(num_segments) else 'cpu'
    batch_shape = torch.as_tensor(batch_shape, dtype=torch.long, device=device)
    assert len(batch_shape.size()) == 1
    num_segments = torch.as_tensor(num_segments, device=device)
    assert len(num_segments.size()) == 0
    indices = torch.arange(start=0, end=num_segments, device=num_segments.device)
    new_tensor = torch.cat([torch.ones_like(batch_shape, dtype=torch.long, device=num_segments.device), num_segments.unsqueeze(dim=0)], dim=0)
    new_shape = [int(x) for x in new_tensor.tolist()]
    indices = indices.view(new_shape)
    multiples = torch.cat([batch_shape, torch.as_tensor([1], device=device)], dim=0)
    indices = indices.repeat(multiples.tolist())
    return IndexMap(indices=indices, num_segments=num_segments, batch_dims=list(batch_shape.size())[0])

Constructs an index map equal to range(num_segments).

Args:
    batch_shape (`torch.Size`):
        Batch shape
    num_segments (`int`):
        Number of segments
    name (`str`, *optional*, defaults to 'range_index_map'):
        Name for the operation. Currently not used

Returns:
    (`IndexMap`): IndexMap of shape batch_shape with elements equal to range(num_segments).


def __invert__(self: EventSetOrNode) -> EventSetOrNode:
    from temporian.core.operators.unary import invert
    return invert(input=self)

Inverts a boolean [`EventSet`][temporian.EventSet] element-wise.

Swaps False <-> True.

Does not work on integers, they should be cast to
[`tp.bool_`][temporian.bool_] beforehand, using
[`EventSet.cast()`][temporian.EventSet.cast].

Example:
    ```python
    >>> a = tp.event_set(
    ...     timestamps=[1, 2],
    ...     features={"M": [1, 5], "N": [1.0, 5.5]},
    ... )
    >>> # Boolean EventSet
    >>> b = a < 2
    >>> b
    indexes: ...
            'M': [ True False]
            'N': [ True False]
    ...

    >>> # Inverted EventSet
    >>> c = ~b
    >>> c
    indexes: ...
            'M': [False True]
            'N': [False True]
    ...

    ```

Returns:
    Inverted EventSet.


class TFOpenAIGPTDoubleHeadsModelOutput(ModelOutput):
    logits: Optional[tf.Tensor] = None
    mc_logits: Optional[tf.Tensor] = None
    hidden_states: Tuple[tf.Tensor] | None = None
    attentions: Tuple[tf.Tensor] | None = None

Base class for outputs of models predicting if two sentences are consecutive or not.

Args:
    logits (`tf.Tensor` of shape `(batch_size, num_choices, sequence_length, config.vocab_size)`):
        Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
    mc_logits (`tf.Tensor` of shape `(batch_size, num_choices)`):
        Prediction scores of the multiple choice classification head (scores for each choice before SoftMax).
    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.


def recipe_dataset(config, auth_write, dataset_dataset, dataset_emails, dataset_groups):
    dataset(config, {'auth': auth_write, 'dataset': dataset_dataset, 'emails': dataset_emails, 'groups': dataset_groups})

Create and permission a dataset in BigQuery.

Args:
  auth_write (authentication) - Credentials used for writing data.
  dataset_dataset (string) - Name of Google BigQuery dataset to create.
  dataset_emails (string_list) - Comma separated emails.
  dataset_groups (string_list) - Comma separated groups.


def max_pool(value, ksize, strides, padding, data_format='NHWC', name=None, input=None):
    value = deprecation.deprecated_argument_lookup('input', input, 'value', value)
    with ops.name_scope(name, 'MaxPool', [value]) as name:
        if data_format is None:
            data_format = 'NHWC'
        channel_index = 1 if data_format.startswith('NC') else 3
        ksize = _get_sequence(ksize, 2, channel_index, 'ksize')
        strides = _get_sequence(strides, 2, channel_index, 'strides')
        if isinstance(padding, (list, tuple)) and data_format == 'NCHW_VECT_C':
            raise ValueError(f"`data_format='NCHW_VECT_C'` is not supported with explicit padding. Received: padding={padding}")
        padding, explicit_paddings = convert_padding(padding)
        if np.isscalar(ksize) and ksize == 0 or (isinstance(ksize, (list, tuple, np.ndarray)) and any((v == 0 for v in ksize))):
            raise ValueError(f'`ksize` cannot be zero. Received: ksize={ksize}')
        return gen_nn_ops.max_pool(value, ksize=ksize, strides=strides, padding=padding, explicit_paddings=explicit_paddings, data_format=data_format, name=name)

Performs the max pooling on the input.

Args:
  value: A 4-D `Tensor` of the format specified by `data_format`.
  ksize: An int or list of `ints` that has length `1`, `2` or `4`.
    The size of the window for each dimension of the input tensor.
  strides: An int or list of `ints` that has length `1`, `2` or `4`.
    The stride of the sliding window for each dimension of the input tensor.
  padding: Either the `string` `"SAME"` or `"VALID"` indicating the type of
    padding algorithm to use, or a list indicating the explicit paddings at
    the start and end of each dimension. When explicit padding is used and
    data_format is `"NHWC"`, this should be in the form `[[0, 0], [pad_top,
    pad_bottom], [pad_left, pad_right], [0, 0]]`. When explicit padding used
    and data_format is `"NCHW"`, this should be in the form `[[0, 0], [0, 0],
    [pad_top, pad_bottom], [pad_left, pad_right]]`. When using explicit
    padding, the size of the paddings cannot be greater than the sliding
    window size.
  data_format: A string. 'NHWC', 'NCHW' and 'NCHW_VECT_C' are supported.
  name: Optional name for the operation.
  input: Alias for value.

Returns:
  A `Tensor` of format specified by `data_format`.
  The max pooled output tensor.


def forward(self, hidden_state, output_hidden_states: bool=False):
    all_hidden_states = []
    embedding = hidden_state
    for mod in self.mixers:
        embedding = mod(embedding)
        if output_hidden_states:
            all_hidden_states.append(embedding)
    if output_hidden_states:
        return (embedding, all_hidden_states)
    else:
        return (embedding, None)

Args:
    hidden_state (`torch.Tensor`): The input tensor.
    output_hidden_states (`bool`, *optional*, defaults to False.):
        Whether to output the hidden states as well.

Returns:
    `torch.Tensor`: The embedding. `list`: List of all hidden states if `output_hidden_states` is set to
    `True`.


def forward(self, hidden_states: torch.Tensor, padding_mask: Optional[torch.Tensor]=False):
    batch_size, sequence_length, hidden_dim = hidden_states.shape
    top_1_mask, router_probs = self.router(hidden_states, padding_mask)
    router_mask = router_probs.bool()
    hidden_states = hidden_states.reshape(batch_size * sequence_length, hidden_dim)
    masked_hidden_states = torch.einsum('bm,be->ebm', hidden_states, router_mask)
    for idx, expert in enumerate(self.experts.values()):
        token_indices = router_mask[:, idx]
        combining_weights = router_probs[token_indices, idx]
        expert_output = expert(masked_hidden_states[idx, token_indices])
        if self.moe_token_dropout > 0:
            if self.training:
                expert_output = self.token_dropout(expert_output)
            else:
                expert_output *= 1 - self.moe_token_dropout
        masked_hidden_states[idx, token_indices] = torch.einsum('b,be->be', combining_weights, expert_output)
    hidden_states = masked_hidden_states.sum(dim=0).reshape(batch_size, sequence_length, hidden_dim)
    top_1_expert_index = torch.argmax(top_1_mask, dim=-1)
    return (hidden_states, (router_probs, top_1_expert_index))

The goal of this forward pass is to have the same number of operation as the equivalent `NllbMoeDenseActDense`
(mlp) layer. This means that all of the hidden states should be processed at most twice ( since we are using a
top_2 gating mechanism). This means that we keep the complexity to O(batch_size x sequence_length x hidden_dim)
instead of O(num_experts x batch_size x sequence_length x hidden_dim).

1- Get the `router_probs` from the `router`. The shape of the `router_mask` is `(batch_size X sequence_length,
num_expert)` and corresponds to the boolean version of the `router_probs`. The inputs are masked using the
`router_mask`.

2- Dispatch the hidden_states to its associated experts. The router probabilities are used to weight the
contribution of each experts when updating the masked hidden states.

Args:
    hidden_states (`torch.Tensor` of shape `(batch_size, sequence_length, hidden_dim)`):
        The hidden states
    padding_mask (`torch.Tensor`, *optional*, defaults to `False`):
        Attention mask. Can be in the causal form or not.

Returns:
    hidden_states (`torch.Tensor` of shape `(batch_size, sequence_length, hidden_dim)`):
        Updated hidden states
    router_logits (`torch.Tensor` of shape `(batch_size, sequence_length, num_experts)`):
        Needed for computing the loss


def _string_to_components(spec=None):
    cached_result = _STRING_TO_COMPONENTS_CACHE.get(spec)
    if cached_result is not None:
        return cached_result
    raw_spec = spec
    job, replica, task, device_type, device_index = (None, None, None, None, None)
    spec = spec or ''
    splits = [x.split(':') for x in spec.split('/')]
    valid_device_types = DeviceSpecV2._get_valid_device_types()
    for y in splits:
        ly = len(y)
        if y:
            if ly == 2 and y[0] == 'job':
                job = y[1]
            elif ly == 2 and y[0] == 'replica':
                replica = y[1]
            elif ly == 2 and y[0] == 'task':
                task = y[1]
            elif (ly == 1 or ly == 2) and y[0].upper() in valid_device_types:
                if device_type is not None:
                    raise ValueError(f'Multiple device types are not allowed while parsing the device spec: {spec}.')
                device_type = y[0].upper()
                if ly == 2 and y[1] != '*':
                    device_index = int(y[1])
            elif ly == 3 and y[0] == 'device':
                if device_type is not None:
                    raise ValueError(f'Multiple device types are not allowed while parsing the device spec: {spec}.')
                device_type = y[1]
                if y[2] != '*':
                    device_index = int(y[2])
            elif ly and y[0] != '':
                raise ValueError(f"Unknown attribute '{y[0]}' is encountered while parsing the device spec: '{spec}'.")
    output = (job, replica, task, device_type, device_index)
    _STRING_TO_COMPONENTS_CACHE[raw_spec] = output
    return output

Stateless portion of device spec string parsing.

Args:
  spec: An optional string specifying a device specification.

Returns:
  The parsed components of `spec`. Note that the result of this function
  must go through attribute setters of DeviceSpec, and should therefore NOT
  be used directly.


def __init__(self, min_value: int=0, max_value: Optional[int]=None):
    super().__init__()
    self._min_value = min_value
    self._max_value = max_value

Constructor.

Args:
  min_value: Min value that is acceptable for the list index.
  max_value: Max value that is acceptable for the list index. If None, there
    is no upper bound for list index.


def lu_solve(lower_upper, perm, rhs, validate_args=False, name=None):
    with ops.name_scope(name or 'lu_solve'):
        lower_upper = ops.convert_to_tensor(lower_upper, dtype_hint=dtypes.float32, name='lower_upper')
        perm = ops.convert_to_tensor(perm, dtype_hint=dtypes.int32, name='perm')
        rhs = ops.convert_to_tensor(rhs, dtype_hint=lower_upper.dtype, name='rhs')
        assertions = _lu_solve_assertions(lower_upper, perm, rhs, validate_args)
        if assertions:
            with ops.control_dependencies(assertions):
                lower_upper = array_ops.identity(lower_upper)
                perm = array_ops.identity(perm)
                rhs = array_ops.identity(rhs)
        if rhs.shape.rank == 2 and perm.shape.rank == 1:
            permuted_rhs = array_ops.gather(rhs, perm, axis=-2)
        else:
            rhs_shape = array_ops.shape(rhs)
            broadcast_batch_shape = array_ops.broadcast_dynamic_shape(rhs_shape[:-2], array_ops.shape(perm)[:-1])
            d, m = (rhs_shape[-2], rhs_shape[-1])
            rhs_broadcast_shape = array_ops.concat([broadcast_batch_shape, [d, m]], axis=0)
            broadcast_rhs = array_ops.broadcast_to(rhs, rhs_broadcast_shape)
            broadcast_rhs = array_ops.reshape(broadcast_rhs, [-1, d, m])
            broadcast_perm = array_ops.broadcast_to(perm, rhs_broadcast_shape[:-1])
            broadcast_perm = array_ops.reshape(broadcast_perm, [-1, d])
            broadcast_batch_size = math_ops.reduce_prod(broadcast_batch_shape)
            broadcast_batch_indices = array_ops.broadcast_to(math_ops.range(broadcast_batch_size)[:, array_ops.newaxis], [broadcast_batch_size, d])
            broadcast_perm = array_ops_stack.stack([broadcast_batch_indices, broadcast_perm], axis=-1)
            permuted_rhs = array_ops.gather_nd(broadcast_rhs, broadcast_perm)
            permuted_rhs = array_ops.reshape(permuted_rhs, rhs_broadcast_shape)
        lower = set_diag(band_part(lower_upper, num_lower=-1, num_upper=0), array_ops.ones(array_ops.shape(lower_upper)[:-1], dtype=lower_upper.dtype))
        return triangular_solve(lower_upper, triangular_solve(lower, permuted_rhs), lower=False)

Solves systems of linear eqns `A X = RHS`, given LU factorizations.

Note: this function does not verify the implied matrix is actually invertible
nor is this condition checked even when `validate_args=True`.

Args:
  lower_upper: `lu` as returned by `tf.linalg.lu`, i.e., if `matmul(P,
    matmul(L, U)) = X` then `lower_upper = L + U - eye`.
  perm: `p` as returned by `tf.linag.lu`, i.e., if `matmul(P, matmul(L, U)) =
    X` then `perm = argmax(P)`.
  rhs: Matrix-shaped float `Tensor` representing targets for which to solve;
    `A X = RHS`. To handle vector cases, use: `lu_solve(..., rhs[...,
      tf.newaxis])[..., 0]`.
  validate_args: Python `bool` indicating whether arguments should be checked
    for correctness. Note: this function does not verify the implied matrix is
      actually invertible, even when `validate_args=True`.
    Default value: `False` (i.e., don't validate arguments).
  name: Python `str` name given to ops managed by this object.
    Default value: `None` (i.e., 'lu_solve').

Returns:
  x: The `X` in `A @ X = RHS`.

#### Examples

```python
import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp

x = [[[1., 2],
      [3, 4]],
     [[7, 8],
      [3, 4]]]
inv_x = tf.linalg.lu_solve(*tf.linalg.lu(x), rhs=tf.eye(2))
tf.assert_near(tf.matrix_inverse(x), inv_x)
# ==> True
```


def set_key(self, structure_prefix, key_line):
    self._empty = False
    key_value = self._remove_structure_prefix(structure_prefix, key_line)
    if '=' in key_value:
        key, value = key_value.split('=', 1)
        self.current_key = key
        if key in self.known_keys:
            self.known_keys[key].append(value)
        else:
            self.unknown_keys[key].append(key_value)

Sets the current key for the instrumentation block.

For unknown keys, the key is added to the value list in order to
better contextualize the value in the output.

Args:
  structure_prefix: string, the structure prefix that was matched
    and that needs to be removed.
  key_line: string, the raw instrumentation output line that contains
    the key-value pair.


def range(*args, **kwargs) -> 'DatasetV2':
    from tensorflow.python.data.ops import range_op
    return range_op._range(*args, **kwargs)

Creates a `Dataset` of a step-separated range of values.

>>> ds = Dataset.range(5)
>>> [a.item() for a in ds.as_numpy_iterator()]
[0, 1, 2, 3, 4]
>>> ds = Dataset.range(2, 5)
>>> [a.item() for a in ds.as_numpy_iterator()]
[2, 3, 4]
>>> ds = Dataset.range(1, 5, 2)
>>> [a.item() for a in ds.as_numpy_iterator()]
[1, 3]
>>> ds = Dataset.range(1, 5, -2)
>>> [a.item() for a in ds.as_numpy_iterator()]
[]
>>> ds = Dataset.range(5, 1)
>>> [a.item() for a in ds.as_numpy_iterator()]
[]
>>> ds = Dataset.range(5, 1, -2)
>>> [a.item() for a in ds.as_numpy_iterator()]
[5, 3]
>>> ds = Dataset.range(2, 5, output_type=tf.int32)
>>> [a.item() for a in ds.as_numpy_iterator()]
[2, 3, 4]
>>> ds = Dataset.range(1, 5, 2, output_type=tf.float32)
>>> [a.item() for a in ds.as_numpy_iterator()]
[1.0, 3.0]

Args:
  *args: follows the same semantics as python's range.
    len(args) == 1 -> start = 0, stop = args[0], step = 1.
    len(args) == 2 -> start = args[0], stop = args[1], step = 1.
    len(args) == 3 -> start = args[0], stop = args[1], step = args[2].
  **kwargs:
    - output_type: Its expected dtype. (Optional, default: `tf.int64`).
    - name: (Optional.) A name for the tf.data operation.

Returns:
  Dataset: A `RangeDataset`.

Raises:
  ValueError: if len(args) == 0.


def from_vision_qformer_text_configs(cls, vision_config: InstructBlipVideoVisionConfig, qformer_config: InstructBlipVideoQFormerConfig, text_config: PretrainedConfig, **kwargs):
    return cls(vision_config=vision_config.to_dict(), qformer_config=qformer_config.to_dict(), text_config=text_config.to_dict(), **kwargs)

Instantiate a [`InstructBlipVideoConfig`] (or a derived class) from a InstructBlipVideo vision model, Q-Former and
language model configurations.

Returns:
    [`InstructBlipVideoConfig`]: An instance of a configuration object


class Buffer(Generic[T]):
    queue: List[T]
    max_size: int
    flusher: Union[FlushFunction, NoReturn]

    def __init__(self, initlist: List[T], max_size: int, flusher: FlushFunction) -> None:
        self.queue = initlist
        self.max_size = max_size
        self.flusher = flusher

    def flush(self, force: bool=False) -> bool | Any:
        """
        Empty and consume queue items, if force or max_size reached

        Args:
            * force: If True, force queue to flush

        Returns:
            * True, if flushed with no errors
            * False, if not flushed
            * Error value from consumer, if flushed with errors
        """
        if force or len(self.queue) > self.max_size:
            result = self.flusher(self.queue)
            self.queue.clear()
            return result or True
        else:
            return False

    def push(self, item: T) -> bool | Any:
        """
        Add item to queue and attempt a flush.

        Args:
            * item: Item to add to queue

        Returns:
            * True, if flushed with no errors
            * False, if not flushed
            * Error value from consumer, if flushed with errors
        """
        self.queue.append(item)
        return self.flush()

Representation of a Buffer (FIFO queue) with the ability to
consume the current queue into a flush function when max_size is reached.

It can queue any list of items, e.g. logs, rows, and API calls.

Args:
    * initlist: Initial list of items
    * max_size: Maximum queue size
    * flusher: Function to be called with list of items


def merge_with(self, other):
    other = as_shape(other)
    if self.dims is None:
        return other
    if other.dims is None:
        return self
    else:
        try:
            self.assert_same_rank(other)
            new_dims = [dim.merge_with(other_dim) for dim, other_dim in zip(self.dims, other.dims)]
            return TensorShape(new_dims)
        except ValueError:
            raise ValueError('Shapes %s and %s are not compatible' % (self, other))

Returns a `TensorShape` combining the information in `self` and `other`.

The dimensions in `self` and `other` are merged element-wise,
according to the rules below:

```python
Dimension(n).merge_with(Dimension(None)) == Dimension(n)
Dimension(None).merge_with(Dimension(n)) == Dimension(n)
Dimension(None).merge_with(Dimension(None)) == Dimension(None)
# raises ValueError for n != m
Dimension(n).merge_with(Dimension(m))
```
>> ts = tf.TensorShape([1,2])
>> ot1 = tf.TensorShape([1,2])
>> ts.merge_with(ot).as_list()
[1,2]

>> ot2 = tf.TensorShape([1,None])
>> ts.merge_with(ot2).as_list()
[1,2]

>> ot3 = tf.TensorShape([None, None])
>> ot3.merge_with(ot2).as_list()
[1, None]

Args:
  other: Another `TensorShape`.

Returns:
  A `TensorShape` containing the combined information of `self` and
  `other`.

Raises:
  ValueError: If `self` and `other` are not compatible.


def _use_temp_cache(self):
    if self._use_tensor_buffer():
        return False
    if self._use_tensor_values_cache():
        return self._parameters.use_temp_cache_var
    else:
        return False

Returns true if the intermediate values should be stacked instead of being stored in a tf.Variable.

Returns:
  A boolean, denoting whether to use a temporary cache or not.