validate_shape=validate_shape, caching_device=caching_device, name=name,
variable_def=None, dtype=dtype, import_scope=None, constraint=None)
def _placeholder_with_default(input, shape, name=None): # pylint: disable=redefined-builtin,unused-argument
x = np.array(input)
if shape is None or any(s is None for s in shape):
return x
return np.reshape(x, shape)
# --- Begin Public Functions --------------------------------------------------
assert_equal = utils.copy_docstring(
tf1.assert_equal,
_assert_equal)
assert_greater = utils.copy_docstring(
tf1.assert_greater,
_assert_greater)
assert_less = utils.copy_docstring(
tf1.assert_less,
_assert_less)
assert_rank = utils.copy_docstring(
tf1.assert_rank,
_assert_rank)
assert_scalar = utils.copy_docstring(
nbatch = int(np.prod(matrix.shape[:-2]))
flat_mat = matrix.reshape(nbatch, dim, dim)
flat_rhs = rhs.reshape(nbatch, dim, rhs.shape[-1])
result = np.empty(flat_rhs.shape)
if np.size(result):
# ValueError: On entry to STRTRS parameter number 7 had an illegal value.
for i, (mat, rh) in enumerate(zip(flat_mat, flat_rhs)):
result[i] = scipy_linalg.solve_triangular(
mat, rh, lower=lower, trans='C' if adjoint else 'N')
return result.reshape(*rhs.shape)
# --- Begin Public Functions --------------------------------------------------
adjoint = utils.copy_docstring(
tf.linalg.adjoint,
lambda matrix, name=None: _matrix_transpose(matrix, conjugate=True))
band_part = utils.copy_docstring(tf.linalg.band_part, _band_part)
cholesky = utils.copy_docstring(
tf.linalg.cholesky, lambda input, name=None: np.linalg.cholesky(input))
cholesky_solve = utils.copy_docstring(tf.linalg.cholesky_solve,
_cholesky_solve)
det = utils.copy_docstring(tf.linalg.det,
lambda input, name=None: np.linalg.det(input))
diag = utils.copy_docstring(tf.linalg.diag, _diag)
from __future__ import division
from __future__ import print_function
# Dependency imports
import numpy as onp
import jax.numpy as np
import tensorflow.compat.v2 as tf
from tensorflow_probability.python.internal.backend.jax import _utils as utils
__all__ = [
'difference',
]
def _difference(a, b, aminusb=True, validate_indices=True):
if not aminusb:
raise NotImplementedError(
'Argument `aminusb != True` is currently unimplemented.')
if not validate_indices:
raise NotImplementedError(
'Argument `validate_indices != True` is currently unimplemented.')
return np.setdiff1d(a, b)
# --- Begin Public Functions --------------------------------------------------
# TODO(b/136555907): Add unit test.
difference = utils.copy_docstring(tf.sets.difference, _difference)
import jax.numpy as np
import tensorflow.compat.v2 as tf
from tensorflow_probability.python.internal.backend.jax import _utils as utils
__all__ = [
'fft',
'fft2d',
'fft3d',
'ifft',
'ifft2d',
'ifft3d',
]
fft = utils.copy_docstring(
tf.signal.fft,
lambda input, name=None: np.fft.fftn(input, axes=[-1]))
fft2d = utils.copy_docstring(
tf.signal.fft2d,
lambda input, name=None: np.fft.fftn(input, axes=[-2, -1]))
fft3d = utils.copy_docstring(
tf.signal.fft3d,
lambda input, name=None: np.fft.fftn(input, axes=[-3, -2, -1]))
ifft = utils.copy_docstring(
tf.signal.ifft,
lambda input, name=None: np.fft.ifftn(input, axes=[-1]))
ifft2d = utils.copy_docstring(
def _argsort(values, axis=-1, direction='ASCENDING', stable=False, name=None): # pylint: disable=unused-argument
"""Numpy implementation of `tf.argsort`."""
if direction == 'ASCENDING':
pass
elif direction == 'DESCENDING':
values = np.negative(values)
else:
raise ValueError('Unrecognized direction: {}.'.format(direction))
return np.argsort(values, axis, kind='stable' if stable else 'quicksort')
def _sort(values, axis=-1, direction='ASCENDING', stable=False, name=None): # pylint: disable=unused-argument
"""Numpy implementation of `tf.sort`."""
if direction == 'ASCENDING':
pass
elif direction == 'DESCENDING':
values = np.negative(values)
else:
raise ValueError('Unrecognized direction: {}.'.format(direction))
result = np.sort(values, axis, kind='stable' if stable else 'quicksort')
if direction == 'DESCENDING':
return np.negative(result)
return result
# --- Begin Public Functions --------------------------------------------------
argsort = utils.copy_docstring(tf.argsort, _argsort)
sort = utils.copy_docstring(tf.sort, _sort)
# 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.
# ============================================================================
"""Experimental Numpy backend."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
# Dependency imports
import tensorflow.compat.v2 as tf
from tensorflow_probability.python.internal.backend.jax import _utils as utils
__all__ = [
'Assert',
'check_numerics',
]
Assert = utils.copy_docstring( # pylint: disable=invalid-name
tf.debugging.Assert,
lambda condition, data, summarize=None, name=None: None)
check_numerics = utils.copy_docstring(tf.debugging.check_numerics,
lambda x, *_, **__: x)
seed=None,
name=None): # pylint: disable=unused-argument
import jax.random as jaxrand # pylint: disable=g-import-not-at-top
if seed is None:
raise ValueError('Must provide PRNGKey to sample in JAX.')
dtype = utils.common_dtype([minval, maxval], dtype_hint=dtype)
maxval = 1 if maxval is None else maxval
shape = _shape([], shape)
return jaxrand.uniform(key=seed,
shape=shape,
dtype=dtype,
minval=minval,
maxval=maxval)
# --- Begin Public Functions --------------------------------------------------
categorical = utils.copy_docstring(
tf.random.categorical, _categorical_jax if JAX_MODE else _categorical)
gamma = utils.copy_docstring(tf.random.gamma,
_gamma_jax if JAX_MODE else _gamma)
normal = utils.copy_docstring(tf.random.normal,
_normal_jax if JAX_MODE else _normal)
poisson = utils.copy_docstring(tf.random.poisson, _poisson)
uniform = utils.copy_docstring(tf.random.uniform,
_uniform_jax if JAX_MODE else _uniform)
# 'case',
# 'cond',
# 'dynamic_partition',
# 'dynamic_stitch',
# 'map_fn',
# 'scan',
]
def _while_loop(
cond,
body,
loop_vars,
shape_invariants=None,
parallel_iterations=10, # pylint: disable=unused-argument
back_prop=True,
swap_memory=False, # pylint: disable=unused-argument
maximum_iterations=None,
name=None): # pylint: disable=unused-argument
i = 0
while (cond(*loop_vars)
and (maximum_iterations is None or i < maximum_iterations)):
loop_vars = body(*loop_vars)
i += 1
return loop_vars
# --- Begin Public Functions --------------------------------------------------
while_loop = utils.copy_docstring(tf.while_loop, _while_loop)
'uint8',
# 'as_string',
# 'bfloat16',
# 'dtypes',
# 'qint16',
# 'qint32',
# 'qint8',
# 'quint16',
# 'quint8',
]
# --- Begin Public Functions --------------------------------------------------
as_dtype = utils.copy_docstring(
tf.as_dtype,
lambda type_value: onp.dtype( # pylint: disable=g-long-lambda
type_value.name if hasattr(type_value, 'name') else type_value).type)
real_dtype = lambda dtype: np.real(np.zeros((0,), dtype=as_dtype(dtype))).dtype
bool = onp.bool # pylint: disable=redefined-builtin
complex128 = np.complex128
complex64 = np.complex64
double = np.double
float16 = np.float16
float32 = np.float32
def _vlog(*_, **__): # pylint: disable=unused-argument
pass
def _warn(*_, **__): # pylint: disable=unused-argument
pass
def _warning(*_, **__): # pylint: disable=unused-argument
pass
# --- Begin Public Functions --------------------------------------------------
TaskLevelStatusMessage = utils.copy_docstring( # pylint: disable=invalid-name
tf1.logging.TaskLevelStatusMessage,
_TaskLevelStatusMessage)
debug = utils.copy_docstring(
tf1.logging.debug,
_debug)
error = utils.copy_docstring(
tf1.logging.error,
_error)
fatal = utils.copy_docstring(
tf1.logging.fatal,
_fatal)
flush = utils.copy_docstring(
# We offer a non SP version just in case SP isn't installed and this
# because logsumexp is often used.
m = _max_mask_non_finite(input_tensor, axis=axis, keepdims=True)
y = input_tensor - m
y = np.exp(y, out=y)
return m + np.log(np.sum(y, axis=_astuple(axis), keepdims=keepdims))
def _top_k(input, k=1, sorted=True, name=None): # pylint: disable=unused-argument,redefined-builtin
raise NotImplementedError
# --- Begin Public Functions --------------------------------------------------
abs = utils.copy_docstring( # pylint: disable=redefined-builtin
tf.math.abs,
lambda x, name=None: np.abs(x))
accumulate_n = utils.copy_docstring(
tf.math.accumulate_n,
lambda inputs, shape=None, tensor_dtype=None, name=None: ( # pylint: disable=g-long-lambda
sum(map(np.array, inputs)).astype(utils.numpy_dtype(tensor_dtype))))
acos = utils.copy_docstring(tf.math.acos, lambda x, name=None: np.arccos(x))
acosh = utils.copy_docstring(tf.math.acosh, lambda x, name=None: np.arccosh(x))
add = utils.copy_docstring(tf.math.add, lambda x, y, name=None: np.add(x, y))
add_n = utils.copy_docstring(
tf.math.add_n, lambda inputs, name=None: sum(map(np.array, inputs)))
# ============================================================================
"""Numpy implementations of sparse functions."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import tensorflow.compat.v2 as tf
from tensorflow_probability.python.internal.backend.jax import _utils as utils
__all__ = [
'to_dense',
]
def _to_dense(sp_input, default_value=0, validate_indices=True, name=None): # pylint: disable=unused-argument
if default_value != 0:
raise NotImplementedError(
'Argument `default_value != 0` is currently unimplemented.')
if not validate_indices:
raise NotImplementedError(
'Argument `validate_indices != True` is currently unimplemented.')
return sp_input
# --- Begin Public Functions --------------------------------------------------
# TODO(b/136555907): Add unit test.
to_dense = utils.copy_docstring(tf.sparse.to_dense, _to_dense)
def _function(
func=None,
input_signature=None,
autograph=True, # pylint: disable=unused-argument
experimental_autograph_options=None, # pylint: disable=unused-argument
experimental_relax_shapes=False,
experimental_compile=None): # pylint: disable=unused-argument
"""Dummy version of `tf.function`."""
# This code path is for the `foo = tf.function(foo, ...)` use case.
if func is not None:
return func
# This code path is for the following use case:
# @tf.function(...)
# def foo(...):
# ...
# This case is equivalent to `foo = tf.function(...)(foo)`.
return lambda inner_function: inner_function
# --- Begin Public Functions --------------------------------------------------
compat = collections.namedtuple(
'compat', 'dimension_value')(lambda dim: None if dim is None else int(dim))
function = utils.copy_docstring(tf.function, _function)
eye = linalg.eye
matmul = linalg.matmul
del collections, tf, utils
# 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.
# ============================================================================
"""Experimental Numpy backend."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
# Dependency imports
import numpy as onp
import jax.numpy as np
import tensorflow.compat.v1 as tf1
import tensorflow.compat.v2 as tf
from tensorflow_probability.python.internal.backend.jax import _utils as utils
__all__ = [
'constant',
]
constant = utils.copy_docstring(
tf1.initializers.constant,
lambda value=0, dtype=tf.dtypes.float32, verify_shape=False: ( # pylint: disable=g-long-lambda
lambda shape, dtype=None, partition_info=None, verify_shape=None: ( # pylint: disable=g-long-lambda
np.ones(shape, dtype=dtype) * value)))
if np.issubdtype(value.dtype, np.integer):
if not np.issubdtype(dtype_hint, np.integer):
return value
return value.astype(dtype_hint)
return np.array(value, dtype=utils.numpy_dtype(dtype or dtype_hint))
def _dimension_value(dimension):
if dimension is None:
return None
return int(dimension)
# --- Begin Public Functions --------------------------------------------------
dimension_value = utils.copy_docstring(tf.compat.dimension_value,
_dimension_value)
class GradientTape(object):
"""tf.GradientTape stub."""
def __init__(self, persistent=False, watch_accessed_variables=True): # pylint: disable=unused-argument
pass
def __enter__(self):
return self
def __exit__(self, typ, value, traceback): # pylint: disable=unused-argument
pass
def watch(self, tensor): # pylint: disable=unused-argument
pass
def _scan( # pylint: disable=unused-argument
fn,
elems,
initializer=None,
parallel_iterations=10,
back_prop=True,
swap_memory=False,
infer_shape=True,
reverse=False,
name=None):
"""Scan implementation."""
out = []
if initializer is None:
arg = elems[0]
elems = elems[1:]
else:
arg = initializer
for x in elems:
arg = fn(arg, x)
out.append(arg)
return np.array(out)
# --- Begin Public Functions --------------------------------------------------
map_fn = utils.copy_docstring(tf.map_fn, _map_fn)
scan = utils.copy_docstring(tf.scan, _scan)
def _transpose(a, perm=None, conjugate=False, name='transpose'): # pylint: disable=unused-argument
x = np.transpose(a, perm)
return np.conjugate(x) if conjugate else x
def _zeros_like(input, dtype=None, name=None): # pylint: disable=redefined-builtin
s = _shape(input)
if isinstance(s, (np.ndarray, onp.generic)):
return np.zeros(s, utils.numpy_dtype(dtype or input.dtype))
return tf.zeros(s, dtype or s.dtype, name)
# --- Begin Public Functions --------------------------------------------------
concat = utils.copy_docstring(
tf.concat,
lambda values, axis, name='concat': ( # pylint: disable=g-long-lambda
np.concatenate([ops.convert_to_tensor(v) for v in values], axis)))
expand_dims = utils.copy_docstring(
tf.expand_dims, lambda input, axis, name=None: np.expand_dims(input, axis))
fill = utils.copy_docstring(
tf.fill,
lambda dims, value, name=None: value * np.ones(dims,
np.array(value).dtype))
gather = utils.copy_docstring(tf.gather, _gather)
gather_nd = utils.copy_docstring(tf.gather_nd, _gather_nd)
reverse = utils.copy_docstring(tf.reverse, _reverse)