# Copyright 2018 Google LLC

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# https://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

"""Types."""

import abc

import datetime

import operator

from typing import Any, Callable, Dict, Iterable, List, MutableMapping, NamedTuple, Optional, TypeVar, Tuple, Union

from apache_beam.utils import shared

import numpy as np

import six

import tensorflow as tf

from tensorflow_model_analysis.proto import metrics_for_slice_pb2

class RaggedTensorValue(

"""

class SparseTensorValue(

"""

class VarLenTensorValue(

cls, values=values, indices=indices, dense_shape=dense_shape)

# pylint: disable=invalid-name

TensorType = Union[tf.Tensor, tf.SparseTensor, tf.RaggedTensor]

TensorOrOperationType = Union[TensorType, tf.Operation]

DictOfTensorType = Dict[str, TensorType]

TensorTypeMaybeDict = Union[TensorType, DictOfTensorType]

DictOfTensorTypeMaybeDict = Dict[str, TensorTypeMaybeDict]

TensorTypeMaybeMultiLevelDict = Union[TensorTypeMaybeDict,

DictOfTensorTypeMaybeDict]

DictOfTypeSpec = Dict[str, tf.TypeSpec]

TypeSpecMaybeDict = Union[tf.TypeSpec, DictOfTypeSpec]

DictOfTypeSpecMaybeDict = Dict[str, TypeSpecMaybeDict]

TypeSpecMaybeMultiLevelDict = Union[TypeSpecMaybeDict, DictOfTypeSpecMaybeDict]

# TODO(b/171992041): Remove tf.compat.v1.SparseTensorValue.

TensorValue = Union[np.ndarray, SparseTensorValue, RaggedTensorValue,

tf.compat.v1.SparseTensorValue]

DictOfTensorValue = Dict[str, TensorValue]

TensorValueMaybeDict = Union[TensorValue, DictOfTensorValue]

DictOfTensorValueMaybeDict = Dict[str, TensorValueMaybeDict]

TensorValueMaybeMultiLevelDict = Union[TensorValueMaybeDict,

DictOfTensorValueMaybeDict]

MetricVariablesType = List[Any]

PrimitiveMetricValueType = Union[float, int, np.number]

ConcreteStructuredMetricValue = TypeVar(

'ConcreteStructuredMetricValue', bound='StructuredMetricValue')

class StructuredMetricValue(abc.ABC):

return self._apply_binary_op(other, operator.pow)

MetricValueType = Union[PrimitiveMetricValueType, np.ndarray,

StructuredMetricValue]

class ValueWithTDistribution(

return float(self.unsampled_value)

# AddMetricsCallback should have the following prototype:

# def add_metrics_callback(features_dict, predictions_dict, labels_dict):

#

# It should create and return a metric_ops dictionary, such that

# metric_ops['metric_name'] = (value_op, update_op), just as in the Trainer.

#

# Note that features_dict, predictions_dict and labels_dict are not

# necessarily dictionaries - they might also be Tensors, depending on what the

# model's eval_input_receiver_fn returns.

# pyformat: disable

AddMetricsCallbackType = Callable[[

TensorTypeMaybeDict, TensorTypeMaybeDict, TensorTypeMaybeDict

], Dict[str, Tuple[TensorType, TensorType]]]

# pyformat: enable

# Type of keys we support for prediction, label and features dictionaries.

FPLKeyType = Union[str, Tuple[str, ...]]

# Dictionary of Tensor values fetched. The dictionary maps original dictionary

# keys => ('node' => value). This type exists for backward compatibility with

# FeaturesPredictionsLabels, new code should use DictOfTensorValue instead.

DictOfFetchedTensorValues = Dict[FPLKeyType, Dict[str, TensorValue]]

FeaturesPredictionsLabels = NamedTuple(

('labels', DictOfFetchedTensorValues)])

# Used in building the model diagnostics table, a MaterializedColumn is a value

# inside of Extracts that will be emitted to file. Note that for strings, the

# values are raw byte strings rather than unicode strings. This is by design, as

# features can have arbitrary bytes values.

MaterializedColumn = NamedTuple(

'MaterializedColumn',

[('name', str),

('value', Union[List[bytes], List[int], List[float], bytes, int, float])])

# Extracts represent data extracted during pipeline processing. In order to

# provide a flexible API, these types are just dicts where the keys are defined

# (reserved for use) by different extractor implementations. For example, the

# FeaturesExtractor stores the data for the features under the key "features",

# LabelsExtractor stores the data for the labels under the key "labels", etc.

Extracts = MutableMapping[str, Any]

# pylint: enable=invalid-name

class ModelLoader:

return with_load_times

class EvalSharedModel(

)

# MaybeMultipleEvalSharedModels represents a parameter that can take on a single

# model or a list of models.

#

# TODO(b/150416505): Deprecate support for dict.

MaybeMultipleEvalSharedModels = Union[EvalSharedModel, List[EvalSharedModel],

Dict[str, EvalSharedModel]]