690 lines
21 KiB
Python
690 lines
21 KiB
Python
# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
|
|
#
|
|
# 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
|
|
#
|
|
# http://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.
|
|
# ==============================================================================
|
|
"""Utility functions for internal use."""
|
|
# pylint: disable=g-direct-tensorflow-import
|
|
|
|
import inspect
|
|
import numbers
|
|
import os
|
|
import re
|
|
import numpy as np
|
|
|
|
from tensorflow.python.framework import dtypes
|
|
from tensorflow.python.framework import indexed_slices
|
|
from tensorflow.python.framework import tensor_util
|
|
from tensorflow.python.ops import array_ops
|
|
from tensorflow.python.ops import control_flow_ops
|
|
from tensorflow.python.ops import math_ops
|
|
from tensorflow.python.ops.numpy_ops import np_arrays
|
|
from tensorflow.python.ops.numpy_ops import np_dtypes
|
|
from tensorflow.python.ops.numpy_ops import np_export
|
|
from tensorflow.python.types import core
|
|
from tensorflow.python.util import nest
|
|
|
|
|
|
def _canonicalize_axis(axis, rank):
|
|
return _canonicalize_axes([axis], rank)[0]
|
|
|
|
|
|
def _canonicalize_axes(axes, rank):
|
|
rank = _maybe_static(rank)
|
|
|
|
if isinstance(rank, core.Tensor):
|
|
canonicalizer = (
|
|
lambda axis: cond(axis < 0, lambda: axis + rank, lambda: axis))
|
|
else:
|
|
canonicalizer = lambda axis: axis + rank if axis < 0 else axis
|
|
|
|
return [canonicalizer(axis) for axis in axes]
|
|
|
|
|
|
def _supports_signature():
|
|
return hasattr(inspect, 'signature')
|
|
|
|
|
|
def _to_tf_type(dtype):
|
|
"""Converts a native python or numpy type to TF DType.
|
|
|
|
Args:
|
|
dtype: Could be a python type, a numpy type or a TF DType.
|
|
|
|
Returns:
|
|
A tensorflow `DType`.
|
|
"""
|
|
return dtypes.as_dtype(dtype)
|
|
|
|
|
|
def _to_numpy_type(dtype):
|
|
"""Converts a native python or TF DType to numpy type.
|
|
|
|
Args:
|
|
dtype: Could be a python type, a numpy type or a TF DType.
|
|
|
|
Returns:
|
|
A NumPy `dtype`.
|
|
"""
|
|
if isinstance(dtype, dtypes.DType):
|
|
return dtype.as_numpy_dtype
|
|
return np.dtype(dtype)
|
|
|
|
|
|
def isscalar(val):
|
|
"""Returns whether `val` is a scalar value or scalar Tensor."""
|
|
if isinstance(val, np_arrays.ndarray):
|
|
val = val.data
|
|
if isinstance(val, core.Tensor):
|
|
ndims = val.shape.ndims
|
|
if ndims is not None:
|
|
return ndims == 0
|
|
else:
|
|
return math_ops.equal(array_ops.rank(val), 0)
|
|
else:
|
|
return np.isscalar(val)
|
|
|
|
|
|
def _has_docstring(f):
|
|
return (f and hasattr(f, '__doc__') and isinstance(f.__doc__, str) and
|
|
f.__doc__)
|
|
|
|
|
|
def _add_blank_line(s):
|
|
if s.endswith('\n'):
|
|
return s + '\n'
|
|
else:
|
|
return s + '\n\n'
|
|
|
|
|
|
def _np_signature(f):
|
|
"""An enhanced inspect.signature that can handle numpy.ufunc."""
|
|
# TODO(wangpeng): consider migrating away from inspect.signature.
|
|
# inspect.signature is supported in Python 3.3.
|
|
if not hasattr(inspect, 'signature'):
|
|
return None
|
|
if f is None:
|
|
return None
|
|
if not isinstance(f, np.ufunc):
|
|
try:
|
|
return inspect.signature(f)
|
|
except ValueError:
|
|
return None
|
|
|
|
def names_from_num(prefix, n):
|
|
if n <= 0:
|
|
return []
|
|
elif n == 1:
|
|
return [prefix]
|
|
else:
|
|
return [prefix + str(i + 1) for i in range(n)]
|
|
|
|
input_names = names_from_num('x', f.nin)
|
|
output_names = names_from_num('out', f.nout)
|
|
keyword_only_params = [('where', True), ('casting', 'same_kind'),
|
|
('order', 'K'), ('dtype', None), ('subok', True),
|
|
('signature', None), ('extobj', None)]
|
|
params = []
|
|
params += [
|
|
inspect.Parameter(name, inspect.Parameter.POSITIONAL_ONLY)
|
|
for name in input_names
|
|
]
|
|
if f.nout > 1:
|
|
params += [
|
|
inspect.Parameter(
|
|
name, inspect.Parameter.POSITIONAL_ONLY, default=None)
|
|
for name in output_names
|
|
]
|
|
params += [
|
|
inspect.Parameter(
|
|
'out',
|
|
inspect.Parameter.POSITIONAL_OR_KEYWORD,
|
|
default=None if f.nout == 1 else (None,) * f.nout)
|
|
]
|
|
params += [
|
|
inspect.Parameter(name, inspect.Parameter.KEYWORD_ONLY, default=default)
|
|
for name, default in keyword_only_params
|
|
]
|
|
return inspect.Signature(params)
|
|
|
|
|
|
# Python 2 doesn't allow keyword-only argument. Python prior to 3.8 doesn't
|
|
# allow positional-only argument. So we conflate positional-only, keyword-only
|
|
# and positional-or-keyword arguments here.
|
|
def _is_compatible_param_kind(a, b):
|
|
|
|
def relax(k):
|
|
if k in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.KEYWORD_ONLY):
|
|
return inspect.Parameter.POSITIONAL_OR_KEYWORD
|
|
return k
|
|
|
|
return relax(a) == relax(b)
|
|
|
|
|
|
def _prepare_np_fun_name_and_fun(np_fun_name, np_fun):
|
|
"""Mutually propagates information between `np_fun_name` and `np_fun`.
|
|
|
|
If one is None and the other is not, we'll try to make the former not None in
|
|
a best effort.
|
|
|
|
Args:
|
|
np_fun_name: name for the np_fun symbol. At least one of np_fun or
|
|
np_fun_name shoud be set.
|
|
np_fun: the numpy function whose docstring will be used.
|
|
|
|
Returns:
|
|
Processed `np_fun_name` and `np_fun`.
|
|
"""
|
|
if np_fun_name is not None:
|
|
assert isinstance(np_fun_name, str)
|
|
if np_fun is not None:
|
|
assert not isinstance(np_fun, str)
|
|
if np_fun is None:
|
|
assert np_fun_name is not None
|
|
try:
|
|
np_fun = getattr(np, str(np_fun_name))
|
|
except AttributeError:
|
|
np_fun = None
|
|
if np_fun_name is None:
|
|
assert np_fun is not None
|
|
np_fun_name = np_fun.__name__
|
|
return np_fun_name, np_fun
|
|
|
|
|
|
def _np_doc_helper(f, np_f, np_fun_name=None, unsupported_params=None,
|
|
link=None):
|
|
"""Helper to get docs."""
|
|
assert np_f or np_fun_name
|
|
if not np_fun_name:
|
|
np_fun_name = np_f.__name__
|
|
doc = 'TensorFlow variant of NumPy\'s `%s`.\n\n' % np_fun_name
|
|
if unsupported_params:
|
|
doc += 'Unsupported arguments: ' + ', '.join(
|
|
'`' + name + '`' for name in unsupported_params) + '.\n\n'
|
|
if _has_docstring(f):
|
|
doc += f.__doc__
|
|
doc = _add_blank_line(doc)
|
|
# TODO(wangpeng): Re-enable the following and choose inlined vs. link to numpy
|
|
# doc according to some global switch.
|
|
doc = _add_np_doc(doc, np_fun_name, np_f, link=link)
|
|
return doc
|
|
|
|
|
|
_np_doc_form = os.getenv('TF_NP_DOC_FORM', '1.16')
|
|
|
|
|
|
def get_np_doc_form():
|
|
"""Gets the form of the original numpy docstrings.
|
|
|
|
Returns:
|
|
See `set_np_doc_form` for the list of valid values.
|
|
"""
|
|
return _np_doc_form
|
|
|
|
|
|
def set_np_doc_form(value):
|
|
r"""Selects the form of the original numpy docstrings.
|
|
|
|
This function sets a global variable that controls how a tf-numpy symbol's
|
|
docstring should refer to the original numpy docstring. If `value` is
|
|
`'inlined'`, the numpy docstring will be verbatim copied into the tf-numpy
|
|
docstring. Otherwise, a link to the original numpy docstring will be
|
|
added. Which numpy version the link points to depends on `value`:
|
|
* `'stable'`: the current stable version;
|
|
* `'dev'`: the current development version;
|
|
* pattern `\d+(\.\d+(\.\d+)?)?`: `value` will be treated as a version number,
|
|
e.g. '1.16'.
|
|
|
|
Args:
|
|
value: the value to set the global variable to.
|
|
"""
|
|
global _np_doc_form
|
|
_np_doc_form = value
|
|
|
|
|
|
class Link:
|
|
|
|
def __init__(self, v):
|
|
self.value = v
|
|
|
|
|
|
class AliasOf:
|
|
|
|
def __init__(self, v):
|
|
self.value = v
|
|
|
|
|
|
class NoLink:
|
|
pass
|
|
|
|
|
|
def generate_link(flag, np_fun_name):
|
|
"""Generates link from numpy function name.
|
|
|
|
Args:
|
|
flag: the flag to control link form. See `set_np_doc_form`.
|
|
np_fun_name: the numpy function name.
|
|
|
|
Returns:
|
|
A string.
|
|
"""
|
|
# Only adds link in this case
|
|
if flag == 'dev':
|
|
template = 'https://numpy.org/devdocs/reference/generated/numpy.%s.html'
|
|
elif flag == 'stable':
|
|
template = (
|
|
'https://numpy.org/doc/stable/reference/generated/numpy.%s.html')
|
|
elif re.match(r'\d+(\.\d+(\.\d+)?)?$', flag):
|
|
# `flag` is the version number
|
|
template = ('https://numpy.org/doc/' + flag +
|
|
'/reference/generated/numpy.%s.html')
|
|
else:
|
|
return None
|
|
return template % np_fun_name
|
|
|
|
|
|
_is_check_link = (os.getenv('TF_NP_CHECK_LINK', 'False') in
|
|
('True', 'true', '1'))
|
|
|
|
|
|
def is_check_link():
|
|
return _is_check_link
|
|
|
|
|
|
def set_check_link(value):
|
|
global _is_check_link
|
|
_is_check_link = value
|
|
|
|
|
|
def _add_np_doc(doc, np_fun_name, np_f, link):
|
|
"""Appends the numpy docstring to `doc`, according to `set_np_doc_form`.
|
|
|
|
See `set_np_doc_form` for how it controls the form of the numpy docstring.
|
|
|
|
Args:
|
|
doc: the docstring to be appended to.
|
|
np_fun_name: the name of the numpy function.
|
|
np_f: (optional) the numpy function.
|
|
link: (optional) which link to use. See `np_doc` for details.
|
|
|
|
Returns:
|
|
`doc` with numpy docstring appended.
|
|
"""
|
|
flag = get_np_doc_form()
|
|
if flag == 'inlined':
|
|
if _has_docstring(np_f):
|
|
doc += 'Documentation for `numpy.%s`:\n\n' % np_fun_name
|
|
# TODO(wangpeng): It looks like code snippets in numpy doc don't work
|
|
# correctly with doctest. Fix that and remove the reformatting of the np_f
|
|
# comment.
|
|
doc += np_f.__doc__.replace('>>>', '>')
|
|
elif isinstance(flag, str):
|
|
if link is None:
|
|
url = generate_link(flag, np_fun_name)
|
|
elif isinstance(link, AliasOf):
|
|
url = generate_link(flag, link.value)
|
|
elif isinstance(link, Link):
|
|
url = link.value
|
|
else:
|
|
url = None
|
|
if url is not None:
|
|
if is_check_link():
|
|
# Imports locally because some builds may not have `requests`
|
|
import requests # pylint: disable=g-import-not-at-top
|
|
r = requests.head(url)
|
|
if r.status_code != 200:
|
|
raise ValueError(
|
|
f'Check link failed at [{url}] with status code {r.status_code}. '
|
|
f'Argument `np_fun_name` is {np_fun_name}.')
|
|
doc += 'See the NumPy documentation for [`numpy.%s`](%s).' % (
|
|
np_fun_name, url)
|
|
return doc
|
|
|
|
|
|
_is_sig_mismatch_an_error = (
|
|
os.getenv('TF_NP_SIG_MISMATCH_IS_ERROR', 'False') in ('True', 'true', '1'))
|
|
|
|
|
|
def is_sig_mismatch_an_error():
|
|
return _is_sig_mismatch_an_error
|
|
|
|
|
|
def set_is_sig_mismatch_an_error(value):
|
|
global _is_sig_mismatch_an_error
|
|
_is_sig_mismatch_an_error = value
|
|
|
|
|
|
def np_doc(np_fun_name, np_fun=None, export=True, unsupported_params=None,
|
|
link=None):
|
|
"""Attachs numpy docstring to a function.
|
|
|
|
Args:
|
|
np_fun_name: name for the np_fun symbol. At least one of np_fun or
|
|
np_fun_name shoud be set.
|
|
np_fun: (optional) the numpy function whose docstring will be used.
|
|
export: whether to export this symbol under module
|
|
`tf.experimental.numpy`. Note that if `export` is `True`, `np_fun` must be
|
|
a function directly under the `numpy` module, not under any submodule of
|
|
`numpy` (e.g. `numpy.random`).
|
|
unsupported_params: (optional) the list of parameters not supported
|
|
by tf.numpy.
|
|
link: (optional) which link to use. If `None`, a default link generated from
|
|
`np_fun_name` will be used. If an instance of `AliasOf`, `link.value` will
|
|
be used in place of `np_fun_name` for the link generation. If an instance
|
|
of `Link`, `link.value` will be used as the whole link. If an instance of
|
|
`NoLink`, no link will be added.
|
|
|
|
Returns:
|
|
A function decorator that attaches the docstring from `np_fun` to the
|
|
decorated function.
|
|
"""
|
|
np_fun_name_orig, np_fun_orig = np_fun_name, np_fun
|
|
np_fun_name, np_fun = _prepare_np_fun_name_and_fun(np_fun_name, np_fun)
|
|
np_sig = _np_signature(np_fun)
|
|
if unsupported_params is None:
|
|
unsupported_params = []
|
|
|
|
def decorator(f):
|
|
"""The decorator."""
|
|
if hasattr(inspect, 'signature') and np_sig is not None:
|
|
try:
|
|
sig = inspect.signature(f)
|
|
except ValueError:
|
|
sig = None
|
|
if sig is not None:
|
|
for name, param in sig.parameters.items():
|
|
np_param = np_sig.parameters.get(name)
|
|
if np_param is None:
|
|
if is_sig_mismatch_an_error():
|
|
raise TypeError(
|
|
f'Cannot find parameter {name} in the numpy function\'s '
|
|
f'signature (which has these parameters: '
|
|
f'{list(np_sig.parameters.keys())}). Argument `np_fun_name` '
|
|
f'is {np_fun_name_orig}. Argument `np_fun` is {np_fun_orig}.')
|
|
else:
|
|
continue
|
|
if (is_sig_mismatch_an_error() and
|
|
not _is_compatible_param_kind(param.kind, np_param.kind)):
|
|
raise TypeError(
|
|
f'Parameter {name} is of kind {param.kind} while in numpy it '
|
|
f'is of kind {np_param.kind}. Argument `np_fun_name` is '
|
|
f'{np_fun_name_orig}. Argument `np_fun` is {np_fun_orig}.')
|
|
has_default = (param.default != inspect.Parameter.empty)
|
|
np_has_default = (np_param.default != inspect.Parameter.empty)
|
|
if is_sig_mismatch_an_error() and has_default != np_has_default:
|
|
raise TypeError(
|
|
'Parameter {} should{} have a default value. Argument '
|
|
'`np_fun_name` is {}. Argument `np_fun` is {}.'.format(
|
|
name, '' if np_has_default else ' not', np_fun_name_orig,
|
|
np_fun_orig))
|
|
for name in np_sig.parameters:
|
|
if name not in sig.parameters:
|
|
unsupported_params.append(name)
|
|
f.__doc__ = _np_doc_helper(
|
|
f, np_fun, np_fun_name=np_fun_name,
|
|
unsupported_params=unsupported_params, link=link)
|
|
if export:
|
|
return np_export.np_export(np_fun_name)(f)
|
|
else:
|
|
return f
|
|
|
|
return decorator
|
|
|
|
|
|
def np_doc_only(np_fun_name, np_fun=None, export=True):
|
|
"""Attachs numpy docstring to a function.
|
|
|
|
This differs from np_doc in that it doesn't check for a match in signature.
|
|
|
|
Args:
|
|
np_fun_name: name for the np_fun symbol. At least one of np_fun or
|
|
np_fun_name shoud be set.
|
|
np_fun: (optional) the numpy function whose docstring will be used.
|
|
export: whether to export this symbol under module
|
|
`tf.experimental.numpy`. Note that if `export` is `True`, `np_f` must be a
|
|
function directly under the `numpy` module, not under any submodule of
|
|
`numpy` (e.g. `numpy.random`).
|
|
|
|
Returns:
|
|
A function decorator that attaches the docstring from `np_fun` to the
|
|
decorated function.
|
|
"""
|
|
np_fun_name, np_fun = _prepare_np_fun_name_and_fun(np_fun_name, np_fun)
|
|
|
|
def decorator(f):
|
|
f.__doc__ = _np_doc_helper(f, np_fun, np_fun_name=np_fun_name)
|
|
if export:
|
|
return np_export.np_export(np_fun_name)(f)
|
|
else:
|
|
return f
|
|
|
|
return decorator
|
|
|
|
|
|
# pylint: disable=g-short-docstring-punctuation,g-no-space-after-docstring-summary,g-docstring-missing-newline,g-doc-return-or-yield,g-doc-args
|
|
@np_doc('finfo')
|
|
def finfo(dtype):
|
|
"""Note that currently it just forwards to the numpy namesake, while
|
|
tensorflow and numpy dtypes may have different properties."""
|
|
return np.finfo(_to_numpy_type(dtype))
|
|
# pylint: enable=g-short-docstring-punctuation,g-no-space-after-docstring-summary,g-docstring-missing-newline,g-doc-return-or-yield,g-doc-args
|
|
|
|
|
|
def _maybe_get_dtype(x):
|
|
"""Returns a numpy type if available from x. Skips if x is numpy.ndarray."""
|
|
# Don't put np.ndarray in this list, because np.result_type looks at the
|
|
# value (not just dtype) of np.ndarray to decide the result type.
|
|
if isinstance(x, numbers.Real):
|
|
return x
|
|
if isinstance(x, (core.Tensor, indexed_slices.IndexedSlices)):
|
|
return _to_numpy_type(x.dtype)
|
|
if isinstance(x, dtypes.DType):
|
|
return x.as_numpy_dtype
|
|
if isinstance(x, (list, tuple)):
|
|
raise ValueError(
|
|
f'Cannot find dtype for type inference from argument `x` of a sequence '
|
|
f'type {type(x)}. For sequences, please call this function on each '
|
|
f'element individually.')
|
|
return x
|
|
|
|
|
|
# Can't use np_doc because np.result_type is a builtin function.
|
|
@np_doc_only('result_type')
|
|
def result_type(*arrays_and_dtypes): # pylint: disable=missing-function-docstring
|
|
arrays_and_dtypes = [
|
|
_maybe_get_dtype(x) for x in nest.flatten(arrays_and_dtypes)
|
|
]
|
|
if not arrays_and_dtypes:
|
|
# If arrays_and_dtypes is an empty list, let numpy decide what the dtype is.
|
|
arrays_and_dtypes = [np.asarray([])]
|
|
return np_dtypes._result_type(*arrays_and_dtypes) # pylint: disable=protected-access
|
|
|
|
|
|
def result_type_unary(a, dtype): # pylint: disable=missing-function-docstring
|
|
"""Find the result type from a single input and a dtype."""
|
|
if dtype:
|
|
# We need to let np_utils.result_type decide the dtype, not tf.zeros_like
|
|
return result_type(dtype)
|
|
|
|
# np_utils.result_type treats string inputs as dtype strings, not as strings.
|
|
# but for unary we want to treat it as a string input.
|
|
if isinstance(a, str):
|
|
return np.unicode_
|
|
elif isinstance(a, bytes):
|
|
return np.bytes_
|
|
|
|
# TF and numpy has different interpretations of Python types such as
|
|
# `float`, so we let `np_utils.result_type` decide.
|
|
return result_type(a)
|
|
|
|
|
|
def _result_type_binary(t1, t2): # pylint: disable=missing-function-docstring
|
|
"""A specialization of result_type for 2 arguments for performance reasons."""
|
|
try:
|
|
return np_dtypes._result_type(_maybe_get_dtype(t1), # pylint: disable=protected-access
|
|
_maybe_get_dtype(t2)) # pylint: disable=protected-access
|
|
except ValueError:
|
|
return result_type(t1, t2)
|
|
|
|
|
|
@np_doc('promote_types')
|
|
def promote_types(type1, type2): # pylint: disable=missing-function-docstring
|
|
type1 = _to_numpy_type(type1)
|
|
type2 = _to_numpy_type(type2)
|
|
return np_dtypes.canonicalize_dtype(np.promote_types(type1, type2))
|
|
|
|
|
|
def tf_broadcast(*args):
|
|
"""Broadcast tensors.
|
|
|
|
Args:
|
|
*args: a list of tensors whose shapes are broadcastable against each other.
|
|
|
|
Returns:
|
|
Tensors broadcasted to the common shape.
|
|
"""
|
|
if len(args) <= 1:
|
|
return args
|
|
sh = array_ops.shape(args[0])
|
|
for arg in args[1:]:
|
|
sh = array_ops.broadcast_dynamic_shape(sh, array_ops.shape(arg))
|
|
return [array_ops.broadcast_to(arg, sh) for arg in args]
|
|
|
|
|
|
# TODO(wangpeng): Move the following functions to a separate file and check for
|
|
# float dtypes in each of them.
|
|
|
|
|
|
def get_static_value(x):
|
|
"""A version of tf.get_static_value that returns None on float dtypes.
|
|
|
|
It returns None on float dtypes in order to avoid breaking gradients.
|
|
|
|
Args:
|
|
x: a tensor.
|
|
|
|
Returns:
|
|
Same as `tf.get_static_value`, except that it returns None when `x` has a
|
|
float dtype.
|
|
"""
|
|
if isinstance(x, core.Tensor) and (x.dtype.is_floating or x.dtype.is_complex):
|
|
return None
|
|
return tensor_util.constant_value(x)
|
|
|
|
|
|
def _maybe_static(x):
|
|
value = get_static_value(x)
|
|
if value is None:
|
|
return x
|
|
else:
|
|
return value
|
|
|
|
|
|
# All the following functions exist becaues get_static_value can't handle
|
|
# their TF counterparts.
|
|
|
|
|
|
def cond(pred, true_fn, false_fn):
|
|
"""A version of tf.cond that tries to evaluate the condition."""
|
|
v = get_static_value(pred)
|
|
if v is None:
|
|
return control_flow_ops.cond(pred, true_fn, false_fn)
|
|
if v:
|
|
return true_fn()
|
|
else:
|
|
return false_fn()
|
|
|
|
|
|
def add(a, b):
|
|
"""A version of tf.add that eagerly evaluates if possible."""
|
|
return _maybe_static(a) + _maybe_static(b)
|
|
|
|
|
|
def subtract(a, b):
|
|
"""A version of tf.subtract that eagerly evaluates if possible."""
|
|
return _maybe_static(a) - _maybe_static(b)
|
|
|
|
|
|
def greater(a, b):
|
|
"""A version of tf.greater that eagerly evaluates if possible."""
|
|
return _maybe_static(a) > _maybe_static(b)
|
|
|
|
|
|
def greater_equal(a, b):
|
|
"""A version of tf.greater_equal that eagerly evaluates if possible."""
|
|
return _maybe_static(a) >= _maybe_static(b)
|
|
|
|
|
|
def less_equal(a, b):
|
|
"""A version of tf.less_equal that eagerly evaluates if possible."""
|
|
return _maybe_static(a) <= _maybe_static(b)
|
|
|
|
|
|
def logical_and(a, b):
|
|
"""A version of tf.logical_and that eagerly evaluates if possible."""
|
|
a_value = get_static_value(a)
|
|
if a_value is not None:
|
|
if np.isscalar(a_value):
|
|
if a_value:
|
|
return _maybe_static(b)
|
|
else:
|
|
return a_value
|
|
else:
|
|
return a_value & _maybe_static(b)
|
|
else:
|
|
return a & _maybe_static(b)
|
|
|
|
|
|
def logical_or(a, b):
|
|
"""A version of tf.logical_or that eagerly evaluates if possible."""
|
|
a_value = get_static_value(a)
|
|
if a_value is not None:
|
|
if np.isscalar(a_value):
|
|
if a_value:
|
|
return a_value
|
|
else:
|
|
return _maybe_static(b)
|
|
else:
|
|
return a_value | _maybe_static(b)
|
|
else:
|
|
return a | _maybe_static(b)
|
|
|
|
|
|
def getitem(a, slice_spec):
|
|
"""A version of __getitem__ that eagerly evaluates if possible."""
|
|
return _maybe_static(a)[slice_spec]
|
|
|
|
|
|
def reduce_all(input_tensor, axis=None, keepdims=False):
|
|
"""A version of tf.reduce_all that eagerly evaluates if possible."""
|
|
v = get_static_value(input_tensor)
|
|
if v is None:
|
|
return math_ops.reduce_all(input_tensor, axis=axis, keepdims=keepdims)
|
|
else:
|
|
return v.all(axis=axis, keepdims=keepdims)
|
|
|
|
|
|
def reduce_any(input_tensor, axis=None, keepdims=False):
|
|
"""A version of tf.reduce_any that eagerly evaluates if possible."""
|
|
v = get_static_value(input_tensor)
|
|
if v is None:
|
|
return math_ops.reduce_any(input_tensor, axis=axis, keepdims=keepdims)
|
|
else:
|
|
return v.any(axis=axis, keepdims=keepdims)
|
|
|
|
|
|
def tf_rank(t):
|
|
r = t.shape.rank
|
|
if r is not None:
|
|
return r
|
|
return array_ops.rank(t)
|