# Copyright 2015 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.

# ==============================================================================

"""Basic arithmetic operators.

See the @{$python/math_ops} guide.

@@add

@@subtract

@@multiply

@@scalar_mul

@@div

@@divide

@@truediv

@@floordiv

@@realdiv

@@truncatediv

@@floor_div

@@truncatemod

@@floormod

@@mod

@@cross

@@add_n

@@abs

@@negative

@@sign

@@reciprocal

@@square

@@round

@@sqrt

@@rsqrt

@@pow

@@exp

@@expm1

@@log

@@log1p

@@sinh

@@cosh

@@asinh

@@acosh

@@atanh

@@ceil

@@floor

@@maximum

@@minimum

@@cos

@@sin

@@lbeta

@@tan

@@acos

@@asin

@@atan

@@atan2

@@lgamma

@@digamma

@@erf

@@erfc

@@squared_difference

@@igamma

@@igammac

@@zeta

@@polygamma

@@betainc

@@rint

@@diag

@@diag_part

@@trace

@@transpose

@@eye

@@matrix_diag

@@matrix_diag_part

@@matrix_band_part

@@matrix_set_diag

@@matrix_transpose

@@matmul

@@norm

@@matrix_determinant

@@matrix_inverse

@@cholesky

@@cholesky_solve

@@matrix_solve

@@matrix_triangular_solve

@@matrix_solve_ls

@@qr

@@self_adjoint_eig

@@self_adjoint_eigvals

@@svd

@@tensordot

@@complex

@@conj

@@imag

@@angle

@@real

@@fft

@@ifft

@@fft2d

@@ifft2d

@@fft3d

@@ifft3d

@@reduce_sum

@@reduce_prod

@@reduce_min

@@reduce_max

@@reduce_mean

@@reduce_all

@@reduce_any

@@reduce_logsumexp

@@count_nonzero

@@accumulate_n

@@einsum

@@bincount

@@cumsum

@@cumprod

@@segment_sum

@@segment_prod

@@segment_min

@@segment_max

@@segment_mean

@@unsorted_segment_sum

@@unsorted_segment_max

@@sparse_segment_sum

@@sparse_segment_mean

@@sparse_segment_sqrt_n

@@argmin

@@argmax

@@setdiff1d

@@where

@@unique

@@edit_distance

@@invert_permutation

"""

from __future__ import absolute_import

from __future__ import division

from __future__ import print_function

import numpy as np

from six.moves import xrange # pylint: disable=redefined-builtin

from tensorflow.python.framework import common_shapes

from tensorflow.python.framework import constant_op

from tensorflow.python.framework import dtypes

from tensorflow.python.framework import graph_util

from tensorflow.python.framework import ops

from tensorflow.python.framework import sparse_tensor

from tensorflow.python.framework import tensor_shape

from tensorflow.python.ops import array_ops

from tensorflow.python.ops import gen_control_flow_ops

from tensorflow.python.ops import gen_data_flow_ops

from tensorflow.python.ops import gen_math_ops

from tensorflow.python.ops import gen_nn_ops

from tensorflow.python.ops import gen_sparse_ops

from tensorflow.python.ops import gen_spectral_ops

from tensorflow.python.ops import gen_state_ops

from tensorflow.python.ops import state_ops

# go/tf-wildcard-import

# pylint: disable=wildcard-import

from tensorflow.python.ops.gen_math_ops import *

# pylint: enable=wildcard-import

from tensorflow.python.util import compat

from tensorflow.python.util.deprecation import deprecated

from tensorflow.python.util.deprecation import deprecated_args

# Aliases for some automatically-generated names.

linspace = gen_math_ops.lin_space

arg_max = deprecated(None, "Use `argmax` instead")(arg_max) # pylint: disable=used-before-assignment

arg_min = deprecated(None, "Use `argmin` instead")(arg_min) # pylint: disable=used-before-assignment

def _set_doc(doc):

def _decorator(func):

func.__doc__ = doc

return func

return _decorator

# pylint: disable=redefined-builtin

@deprecated_args(None, "Use the `axis` argument instead", "dimension")

@_set_doc(gen_math_ops.arg_max.__doc__

.replace("dimensions", "axes")

.replace("dimension", "axis"))

def argmax(input,

axis=None,

name=None,

dimension=None,

output_type=dtypes.int64):

if dimension is not None:

if axis is not None:

raise ValueError("Cannot specify both 'axis' and 'dimension'")

axis = dimension

elif axis is None:

axis = 0

return gen_math_ops.arg_max(input, axis, name=name, output_type=output_type)

@deprecated_args(None, "Use the `axis` argument instead", "dimension")

@_set_doc(gen_math_ops.arg_min.__doc__

.replace("dimensions", "axes")

.replace("dimension", "axis"))

def argmin(input,

axis=None,

name=None,

dimension=None,

output_type=dtypes.int64):

if dimension is not None:

if axis is not None:

raise ValueError("Cannot specify both 'axis' and 'dimension'")

axis = dimension

elif axis is None:

axis = 0

return gen_math_ops.arg_min(input, axis, name=name, output_type=output_type)

# pylint: enable=redefined-builtin

# pylint: disable=anomalous-backslash-in-string,protected-access

# pylint: disable=g-docstring-has-escape

def abs(x, name=None):

r"""Computes the absolute value of a tensor.

Given a tensor `x` of complex numbers, this operation returns a tensor of type

`float32` or `float64` that is the absolute value of each element in `x`. All

elements in `x` must be complex numbers of the form \\(a + bj\\). The

absolute value is computed as \\( \sqrt{a^2 + b^2}\\). For example:

```python

x = tf.constant([[-2.25 + 4.75j], [-3.25 + 5.75j]])

tf.abs(x) # [5.25594902, 6.60492229]

```

Args:

x: A `Tensor` or `SparseTensor` of type `float32`, `float64`, `int32`,

`int64`, `complex64` or `complex128`.

name: A name for the operation (optional).

Returns:

A `Tensor` or `SparseTensor` the same size and type as `x` with absolute

values.

Note, for `complex64` or `complex128' input, the returned `Tensor` will be

of type `float32` or `float64`, respectively.

"""

with ops.name_scope(name, "Abs", [x]) as name:

if isinstance(x, sparse_tensor.SparseTensor):

if x.values.dtype in (dtypes.complex64, dtypes.complex128):

x_abs = gen_math_ops._complex_abs(

x.values, Tout=x.values.dtype.real_dtype, name=name)

return sparse_tensor.SparseTensor(

indices=x.indices, values=x_abs, dense_shape=x.dense_shape)

x_abs = gen_math_ops._abs(x.values, name=name)

return sparse_tensor.SparseTensor(

indices=x.indices, values=x_abs, dense_shape=x.dense_shape)

else:

x = ops.convert_to_tensor(x, name="x")

if x.dtype in (dtypes.complex64, dtypes.complex128):

return gen_math_ops._complex_abs(x, Tout=x.dtype.real_dtype, name=name)

return gen_math_ops._abs(x, name=name)

# pylint: enable=g-docstring-has-escape

# pylint: disable=redefined-builtin

def _bucketize(input, boundaries, name=None):

return gen_math_ops._bucketize(input=input, boundaries=boundaries, name=name)

# pylint: enable=redefined-builtin

class DivideDelegateWithName(object):

"""Use Python2/Python3 division delegation to implement divide for tensors."""

def __init__(self, x, name):

"""Construct DivideDelegateWithName.

Args:

x: Tensor to use as left operand in operator overloads

name: The name that is preferred for the op created.

"""

self.x = x

self.name = name

def __truediv__(self, y):

return _truediv_python3(self.x, y, self.name)

def __floordiv__(self, y):

return floordiv(self.x, y, self.name)

def __div__(self, y):

return _div_python2(self.x, y, self.name)

def divide(x, y, name=None):

"""Computes Python style division of `x` by `y`."""

if name is not None:

# Cannot use tensors operator overload, because it has no way to track

# override names. Use a dummy class to track the runtime division behavior

return DivideDelegateWithName(x, name) / y

else:

return x / y

def multiply(x, y, name=None):

return gen_math_ops._mul(x, y, name)

multiply.__doc__ = gen_math_ops._mul.__doc__.replace("Mul", "`tf.multiply`")

# TODO(aselle): put deprecation in after another round of global code changes

@deprecated(

"2016-12-30",

"`tf.mul(x, y)` is deprecated, please use `tf.multiply(x, y)` or `x * y`")

def _mul(x, y, name=None):

return gen_math_ops._mul(x, y, name)

_mul.__doc__ = (gen_math_ops._mul.__doc__ +

("" if _mul.__doc__ is None else _mul.__doc__))

def subtract(x, y, name=None):

return gen_math_ops._sub(x, y, name)

subtract.__doc__ = gen_math_ops._sub.__doc__.replace("`Sub`", "`tf.subtract`")

# TODO(aselle): put deprecation in after another round of global code changes

@deprecated(

"2016-12-30",

"`tf.sub(x, y)` is deprecated, please use `tf.subtract(x, y)` or `x - y`")

def _sub(x, y, name=None):

return gen_math_ops._sub(x, y, name)

_sub.__doc__ = (gen_math_ops._sub.__doc__ +

("" if _sub.__doc__ is None else _sub.__doc__))

# pylint: disable=g-docstring-has-escape

def negative(x, name=None):

"""Computes numerical negative value element-wise.

I.e., \\(y = -x\\).

Args:

x: A `Tensor` or `SparseTensor`. Must be one of the following types: `half`,

`float32`, `float64`, `int32`, `int64`, `complex64`, `complex128`.

name: A name for the operation (optional).

Returns:

A `Tensor` or `SparseTensor`, respectively. Has the same type as `x`.

"""

with ops.name_scope(name, "Neg", [x]) as name:

if isinstance(x, sparse_tensor.SparseTensor):

x_neg = gen_math_ops._neg(x.values, name=name)

return sparse_tensor.SparseTensor(

indices=x.indices, values=x_neg, dense_shape=x.dense_shape)

else:

return gen_math_ops._neg(x, name=name)

# pylint: enable=g-docstring-has-escape

# pylint: disable=g-docstring-has-escape

@deprecated("2016-12-30",

"`tf.neg(x)` is deprecated, please use `tf.negative(x)` or `-x`")

def _neg(x, name=None):

"""Computes numerical negative value element-wise.

I.e., \\(y = -x\\).

Args:

x: A `Tensor` or `SparseTensor`. Must be one of the following types: `half`,

`float32`, `float64`, `int32`, `int64`, `complex64`, `complex128`.

name: A name for the operation (optional).

Returns:

A `Tensor` or `SparseTensor`, respectively. Has the same type as `x`.

"""

return negative(x, name)

# pylint: enable=g-docstring-has-escape

def sign(x, name=None):

"""Returns an element-wise indication of the sign of a number.

`y = sign(x) = -1` if `x < 0`; 0 if `x == 0` or `tf.is_nan(x)`; 1 if `x > 0`.

Zero is returned for NaN inputs.

For complex numbers, `y = sign(x) = x / |x|` if `x != 0`, otherwise `y = 0`.

Args:

x: A `Tensor` or `SparseTensor`. Must be one of the following types: `half`,

`float32`, `float64`, `int32`, `int64`, `complex64`, `complex128`.

name: A name for the operation (optional).

Returns:

A `Tensor` or `SparseTensor`, respectively. Has the same type as `x`.

@compatibility(numpy)

Equivalent to numpy.sign except for the behavior for input values of NaN.

@end_compatibility

"""

with ops.name_scope(name, "Sign", [x]) as name:

if isinstance(x, sparse_tensor.SparseTensor):

x_sign = gen_math_ops.sign(x.values, name=name)

return sparse_tensor.SparseTensor(

indices=x.indices, values=x_sign, dense_shape=x.dense_shape)

else:

return gen_math_ops.sign(x, name=name)

def square(x, name=None):

r"""Computes square of x element-wise.

I.e., \\(y = x * x = x^2\\).

Args:

x: A `Tensor` or `SparseTensor`. Must be one of the following types: `half`,

`float32`, `float64`, `int32`, `int64`, `complex64`, `complex128`.

name: A name for the operation (optional).

Returns:

A `Tensor` or `SparseTensor`. Has the same type as `x`.

"""

with ops.name_scope(name, "Square", [x]) as name:

if isinstance(x, sparse_tensor.SparseTensor):

x_square = gen_math_ops.square(x.values, name=name)

return sparse_tensor.SparseTensor(

indices=x.indices, values=x_square, dense_shape=x.dense_shape)

else:

return gen_math_ops.square(x, name=name)

def sqrt(x, name=None):

r"""Computes square root of x element-wise.

I.e., \\(y = \sqrt{x} = x^{1/2}\\).

Args:

x: A `Tensor` or `SparseTensor`. Must be one of the following types: `half`,

`float32`, `float64`, `complex64`, `complex128`.

name: A name for the operation (optional).

Returns:

A `Tensor` or `SparseTensor`, respectively. Has the same type as `x`.

"""

with ops.name_scope(name, "Sqrt", [x]) as name:

if isinstance(x, sparse_tensor.SparseTensor):

x_sqrt = gen_math_ops.sqrt(x.values, name=name)

return sparse_tensor.SparseTensor(

indices=x.indices, values=x_sqrt, dense_shape=x.dense_shape)

else:

return gen_math_ops.sqrt(x, name=name)

def erf(x, name=None):

"""Computes the Gauss error function of `x` element-wise.

Args:

x: A `Tensor` of `SparseTensor`. Must be one of the following types: `half`,

`float32`, `float64`.

name: A name for the operation (optional).

Returns:

A `Tensor` or `SparseTensor`, respectively. Has the same type as `x`.

"""

with ops.name_scope(name, "Erf", [x]) as name:

if isinstance(x, sparse_tensor.SparseTensor):

x_erf = gen_math_ops.erf(x.values, name=name)

return sparse_tensor.SparseTensor(

indices=x.indices, values=x_erf, dense_shape=x.dense_shape)

else:

return gen_math_ops.erf(x, name=name)

def scalar_mul(scalar, x):

"""Multiplies a scalar times a `Tensor` or `IndexedSlices` object.

Intended for use in gradient code which might deal with `IndexedSlices`

objects, which are easy to multiply by a scalar but more expensive to

multiply with arbitrary tensors.

Args:

scalar: A 0-D scalar `Tensor`. Must have known shape.

x: A `Tensor` or `IndexedSlices` to be scaled.

Returns:

`scalar * x` of the same type (`Tensor` or `IndexedSlices`) as `x`.

Raises:

ValueError: if scalar is not a 0-D `scalar`.

"""

scalar = ops.convert_to_tensor(

scalar, dtype=x.dtype.base_dtype, name="scalar")

shape = scalar.get_shape()

if shape.ndims == 0:

if isinstance(x, ops.IndexedSlices):

return ops.IndexedSlices(scalar * x.values, x.indices, x.dense_shape)

else:

return scalar * x

else:

raise ValueError("Only scalar multiply works, got shape %s" % shape)

def pow(x, y, name=None):

r"""Computes the power of one value to another.

Given a tensor `x` and a tensor `y`, this operation computes \\(x^y\\) for

corresponding elements in `x` and `y`. For example:

```python

x = tf.constant([[2, 2], [3, 3]])

y = tf.constant([[8, 16], [2, 3]])

tf.pow(x, y) # [[256, 65536], [9, 27]]

```

Args:

x: A `Tensor` of type `float32`, `float64`, `int32`, `int64`, `complex64`,

or `complex128`.

y: A `Tensor` of type `float32`, `float64`, `int32`, `int64`, `complex64`,

or `complex128`.

name: A name for the operation (optional).

Returns:

A `Tensor`.

"""

with ops.name_scope(name, "Pow", [x]) as name:

return gen_math_ops._pow(x, y, name=name)

# pylint: disable=redefined-builtin,redefined-outer-name

def complex(real, imag, name=None):

r"""Converts two real numbers to a complex number.

Given a tensor `real` representing the real part of a complex number, and a

tensor `imag` representing the imaginary part of a complex number, this

operation returns complex numbers elementwise of the form \\(a + bj\\), where

*a* represents the `real` part and *b* represents the `imag` part.

The input tensors `real` and `imag` must have the same shape.

For example:

```python

real = tf.constant([2.25, 3.25])

imag = tf.constant([4.75, 5.75])

tf.complex(real, imag) # [[2.25 + 4.75j], [3.25 + 5.75j]]

```

Args:

real: A `Tensor`. Must be one of the following types: `float32`,

`float64`.

imag: A `Tensor`. Must have the same type as `real`.

name: A name for the operation (optional).

Returns:

A `Tensor` of type `complex64` or `complex128`.

"""

real = ops.convert_to_tensor(real, name="real")

imag = ops.convert_to_tensor(imag, name="imag")

with ops.name_scope(name, "Complex", [real, imag]) as name:

input_types = (real.dtype, imag.dtype)

if input_types == (dtypes.float64, dtypes.float64):

Tout = dtypes.complex128

elif input_types == (dtypes.float32, dtypes.float32):

Tout = dtypes.complex64

else:

raise TypeError("real and imag have incorrect types: "

"{} {}".format(real.dtype.name, imag.dtype.name))

return gen_math_ops._complex(real, imag, Tout=Tout, name=name)

def real(input, name=None):

r"""Returns the real part of a complex number.

Given a tensor `input` of complex numbers, this operation returns a tensor of

type `float32` or `float64` that is the real part of each element in `input`.

All elements in `input` must be complex numbers of the form \\(a + bj\\),

where *a* is the real part returned by this operation and *b* is the

imaginary part.

For example:

```python

x = tf.constant([-2.25 + 4.75j, 3.25 + 5.75j])

tf.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`.

"""

with ops.name_scope(name, "Real", [input]) as name:

real_dtype = input.dtype.real_dtype

if input.dtype.base_dtype == real_dtype:

return input

return gen_math_ops.real(input, Tout=real_dtype, name=name)

def imag(input, name=None):

r"""Returns the imaginary part of a complex number.

Given a tensor `input` of complex numbers, this operation returns a tensor of

type `float` that is the argument of each element in `input`. All elements in

`input` must be complex numbers of the form \\(a + bj\\), where *a*

is the real part and *b* is the imaginary part returned by the operation.

For example:

```python

x = tf.constant([-2.25 + 4.75j, 3.25 + 5.75j])

tf.imag(x) # [4.75, 5.75]

```

Args:

input: A `Tensor`. Must be one of the following types: `complex64`,

`complex128`.

name: A name for the operation (optional).

Returns:

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

"""

with ops.name_scope(name, "Imag", [input]) as name:

return gen_math_ops.imag(input, Tout=input.dtype.real_dtype, name=name)

def angle(input, name=None):

r"""Returns the argument of a complex number.

Given a tensor `input` of complex numbers, this operation returns a tensor of

type `float32` or `float64` that is the argument of each element in `input`.

All elements in `input` must be complex numbers of the form \\(a + bj\\),

where *a* is the real part and *b* is the imaginary part.

The argument returned by this function is of the form \\(atan2(b, a)\\).

For example:

```

# tensor 'input' is [-2.25 + 4.75j, 3.25 + 5.75j]

tf.angle(input) ==> [2.0132, 1.056]

```

Args:

input: A `Tensor`. Must be one of the following types: `complex64`,

`complex128`.

name: A name for the operation (optional).

Returns:

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

"""

with ops.name_scope(name, "Angle", [input]) as name:

return gen_math_ops.angle(input, Tout=input.dtype.real_dtype, name=name)

# pylint: enable=redefined-outer-name,redefined-builtin

def round(x, name=None):

"""Rounds the values of a tensor to the nearest integer, element-wise.

Rounds half to even. Also known as bankers rounding. If you want to round

according to the current system rounding mode use tf::cint.

For example:

```python

x = tf.constant([0.9, 2.5, 2.3, 1.5, -4.5])

tf.round(x) # [ 1.0, 2.0, 2.0, 2.0, -4.0 ]

```

Args:

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

name: A name for the operation (optional).

Returns:

A `Tensor` of same shape and type as `x`.

"""

x = ops.convert_to_tensor(x, name="x")

if x.dtype.is_integer:

return x

else:

return gen_math_ops.round(x, name=name)

def cast(x, dtype, name=None):

"""Casts a tensor to a new type.

The operation casts `x` (in case of `Tensor`) or `x.values`

(in case of `SparseTensor`) to `dtype`.

For example:

```python

x = tf.constant([1.8, 2.2], dtype=tf.float32)

tf.cast(x, tf.int32) # [1, 2], dtype=tf.int32

```

Args:

x: A `Tensor` or `SparseTensor`.

dtype: The destination type.

name: A name for the operation (optional).

Returns:

A `Tensor` or `SparseTensor` with same shape as `x`.

Raises:

TypeError: If `x` cannot be cast to the `dtype`.

"""

base_type = dtypes.as_dtype(dtype).base_dtype

with ops.name_scope(name, "Cast", [x]) as name:

if isinstance(x, sparse_tensor.SparseTensor):

values_cast = cast(x.values, base_type, name=name)

return sparse_tensor.SparseTensor(x.indices, values_cast, x.dense_shape)

else:

# TODO(touts): Handle what Josh said.

#

# Could return ops.convert_to_tensor(x, dtype=dtype, ...) here, but that

# allows some conversions that cast() can't do, e.g. casting numbers to

# strings.

x = ops.convert_to_tensor(x, name="x")

if x.dtype.base_dtype == base_type:

return x

return gen_math_ops.cast(x, base_type, name=name)

def saturate_cast(value, dtype, name=None):

"""Performs a safe saturating cast of `value` to `dtype`.

This function casts the input to `dtype` without applying any scaling. If

there is a danger that values would over or underflow in the cast, this op

applies the appropriate clamping before the cast.

Args:

value: A `Tensor`.

dtype: The desired output `DType`.

name: A name for the operation (optional).

Returns:

`value` safely cast to `dtype`.

"""

# When casting to a type with smaller representable range, clamp.

# Note that this covers casting to unsigned types as well.

with ops.name_scope(name, "saturate_cast", [value]) as name:

value = ops.convert_to_tensor(value, name="value")

dtype = dtypes.as_dtype(dtype).base_dtype

if value.dtype.min < dtype.min:

value = gen_math_ops.maximum(value,

ops.convert_to_tensor(

dtype.min, dtype=value.dtype,

name="min"))

if value.dtype.max > dtype.max:

value = gen_math_ops.minimum(value,

ops.convert_to_tensor(

dtype.max, dtype=value.dtype,

name="max"))

return cast(value, dtype, name=name)

def to_float(x, name="ToFloat"):

"""Casts a tensor to type `float32`.

Args:

x: A `Tensor` or `SparseTensor`.

name: A name for the operation (optional).

Returns:

A `Tensor` or `SparseTensor` with same shape as `x` with type `float32`.

Raises:

TypeError: If `x` cannot be cast to the `float32`.

"""

return cast(x, dtypes.float32, name=name)

def to_double(x, name="ToDouble"):

"""Casts a tensor to type `float64`.

Args:

x: A `Tensor` or `SparseTensor`.

name: A name for the operation (optional).

Returns:

A `Tensor` or `SparseTensor` with same shape as `x` with type `float64`.

Raises:

TypeError: If `x` cannot be cast to the `float64`.

"""

return cast(x, dtypes.float64, name=name)

def to_int32(x, name="ToInt32"):

"""Casts a tensor to type `int32`.

Args:

x: A `Tensor` or `SparseTensor`.

name: A name for the operation (optional).

Returns:

A `Tensor` or `SparseTensor` with same shape as `x` with type `int32`.

Raises:

TypeError: If `x` cannot be cast to the `int32`.

"""

return cast(x, dtypes.int32, name=name)

def to_int64(x, name="ToInt64"):

"""Casts a tensor to type `int64`.

Args:

x: A `Tensor` or `SparseTensor`.

name: A name for the operation (optional).

Returns:

A `Tensor` or `SparseTensor` with same shape as `x` with type `int64`.

Raises:

TypeError: If `x` cannot be cast to the `int64`.

"""

return cast(x, dtypes.int64, name=name)

def to_bfloat16(x, name="ToBFloat16"):

"""Casts a tensor to type `bfloat16`.

Args:

x: A `Tensor` or `SparseTensor`.

name: A name for the operation (optional).

Returns:

A `Tensor` or `SparseTensor` with same shape as `x` with type `bfloat16`.

Raises:

TypeError: If `x` cannot be cast to the `bfloat16`.

"""

return cast(x, dtypes.bfloat16, name=name)

ops.Tensor._override_operator("__neg__", gen_math_ops._neg)

ops.Tensor._override_operator("__abs__", abs)

# __invert__ corresponds to the ~ operator. Here we follow the numpy convention

# ~ marks an elementwise bit-wise inverse. This is only implemented for boolean

# tensors and will throw a TypeError if used on nonboolean arrays

ops.Tensor._override_operator("__invert__", gen_math_ops.logical_not)

def _OverrideBinaryOperatorHelper(func, op_name, clazz_object=ops.Tensor):

"""Register operators with different tensor and scalar versions.

If `clazz_object` is `SparseTensor`, assumes `func` takes `(sp_indices,

sp_values, sp_shape, dense)` and outputs `(new_sp_values)`.

Args:

func: the operator

op_name: name of the operator being overridden

clazz_object: class to override for. Either `Tensor` or `SparseTensor`.

"""

def binary_op_wrapper(x, y):

with ops.name_scope(None, op_name, [x, y]) as name:

if not isinstance(y, sparse_tensor.SparseTensor):

try:

y = ops.convert_to_tensor(y, dtype=x.dtype.base_dtype, name="y")

except TypeError:

# If the RHS is not a tensor, it might be a tensor aware object

# that can implement the operator with knowledge of itself

# and the tensor.

if hasattr(type(y), "__r%s__" % op_name):

return NotImplemented

else:

raise

return func(x, y, name=name)

def binary_op_wrapper_sparse(sp_x, y):

with ops.name_scope(None, op_name, [sp_x, y]) as name:

y = ops.convert_to_tensor(y, dtype=sp_x.dtype.base_dtype, name="y")

return sparse_tensor.SparseTensor(sp_x.indices,

func(

sp_x.indices,

sp_x.values,

sp_x.dense_shape,

y,

name=name), sp_x.dense_shape)

def r_binary_op_wrapper(y, x):

with ops.name_scope(None, op_name, [x, y]) as name:

x = ops.convert_to_tensor(x, dtype=y.dtype.base_dtype, name="x")

return func(x, y, name=name)

# Propagate func.__doc__ to the wrappers

try:

doc = func.__doc__

except AttributeError:

doc = None

binary_op_wrapper.__doc__ = doc

r_binary_op_wrapper.__doc__ = doc

binary_op_wrapper_sparse.__doc__ = doc

if clazz_object is ops.Tensor:

clazz_object._override_operator("__%s__" % op_name, binary_op_wrapper)

del binary_op_wrapper

clazz_object._override_operator("__r%s__" % op_name, r_binary_op_wrapper)

del r_binary_op_wrapper

else:

clazz_object._override_operator("__%s__" % op_name,

binary_op_wrapper_sparse)

del binary_op_wrapper_sparse

# Conversion table for __truediv__. None entries mean no conversion required.

_TRUEDIV_TABLE = {

dtypes.uint8: dtypes.float32,

dtypes.int8: dtypes.float32,

dtypes.uint16: dtypes.float32,

dtypes.int16: dtypes.float32,

dtypes.int32: dtypes.float64,

dtypes.int64: dtypes.float64,

dtypes.float16: None,

dtypes.float32: None,

dtypes.float64: None,

dtypes.complex64: None,

dtypes.complex128: None,

}

# NOTE: the support of "sparse (true)div dense" is currently not baked in into

# "tf.(true_)div()". Until such an API decision is made, the supported usage is

# to explicitly use the "/" operator to invoke either truediv or div.

def _sparse_dense_truediv(sp_indices, sp_values, sp_shape, y, name=None):

"""Internal helper function for 'sp_t / dense_t'."""

with ops.name_scope(name, "truediv", [sp_indices, sp_values, sp_shape,

y]) as name:

sp_values = ops.convert_to_tensor(sp_values, name="sp_values")

y = ops.convert_to_tensor(y, name="y")

x_dtype = sp_values.dtype.base_dtype

y_dtype = y.dtype.base_dtype

if x_dtype != y_dtype:

raise TypeError("x and y must have the same dtype, got %r != %r" %

(x_dtype, y_dtype))

try:

dtype = _TRUEDIV_TABLE[x_dtype]

except KeyError:

raise TypeError("Invalid dtype %r in __truediv__" % x_dtype)

if dtype is not None:

sp_values = cast(sp_values, dtype)

y = cast(y, dtype)

return gen_sparse_ops.sparse_dense_cwise_div(

sp_indices, sp_values, sp_shape, y, name=name)

def _truediv_python3(x, y, name=None):

with ops.name_scope(name, "truediv", [x, y]) as name:

x = ops.convert_to_tensor(x, name="x")

y = ops.convert_to_tensor(y, name="y")

x_dtype = x.dtype.base_dtype

y_dtype = y.dtype.base_dtype

if x_dtype != y_dtype:

raise TypeError("x and y must have the same dtype, got %r != %r" %

(x_dtype, y_dtype))

try:

dtype = _TRUEDIV_TABLE[x_dtype]

except KeyError:

raise TypeError("Invalid dtype %r in __truediv__" % x_dtype)

if dtype is not None:

x = cast(x, dtype)

y = cast(y, dtype)

return gen_math_ops._real_div(x, y, name=name)

def _div_python2(x, y, name=None):

"""Divide two values using Python 2 semantics. Used for Tensor.__div__.

Args:

x: `Tensor` numerator of real numeric type.

y: `Tensor` denominator of real numeric type.

name: A name for the operation (optional).

Returns:

`x / y` returns the quotient of x and y.