def _bin_op(tf_fun, a, b, promote=True):
if promote:
a, b = array_creation.promote_args_types(a, b)
else:
a = array_creation.asarray(a)
b = array_creation.asarray(b)
return utils.tensor_to_ndarray(tf_fun(a.data, b.data))
def kron(a, b):
a, b = array_creation.promote_args_types(a, b)
ndim = __builtins__["max"](a.ndim, b.ndim)
if a.ndim < ndim:
a = array_methods.reshape(a, _pad_left_to(ndim, a.shape))
if b.ndim < ndim:
b = array_methods.reshape(b, _pad_left_to(ndim, b.shape))
a_reshaped = array_methods.reshape(a, [i for d in a.shape for i in (d, 1)])
b_reshaped = array_methods.reshape(b, [i for d in b.shape for i in (1, d)])
out_shape = tuple(np.multiply(a.shape, b.shape))
return array_methods.reshape(a_reshaped * b_reshaped, out_shape)
def where(condition, x, y):
"""Return an array with elements from `x` or `y`, depending on condition.
Args:
condition: array_like, bool. Where True, yield `x`, otherwise yield `y`.
x: see below.
y: array_like, optional. Values from which to choose. `x`, `y` and
`condition` need to be broadcastable to some shape.
Returns:
An array.
"""
condition = array_creation.asarray(condition, dtype=np.bool_)
x, y = array_creation.promote_args_types(x, y)
return utils.tensor_to_ndarray(tf.where(condition.data, x.data, y.data))
def dot(a, b):
"""The dot product of two arrays. See numpy.dot for more details.
This relies on `tf.tensordot` which does not support types int64 and float64.
So arrays of those types are "unsafely" cast to int32 and float32.
Args:
a: array_like. Could be an ndarray, a Tensor or any object that can
be converted to a Tensor using `tf.convert_to_tensor`.
b: array_like. Could be an ndarray, a Tensor or any object that can
be converted to a Tensor using `tf.convert_to_tensor`.
Returns:
An ndarray.
"""
a, b = array_creation.promote_args_types(a, b)
if utils.isscalar(a) or utils.isscalar(b):
a = utils.tensor_to_ndarray(tf.expand_dims(a.data, -1))
b = utils.tensor_to_ndarray(tf.expand_dims(b.data, -1))
a_axis = b_axis = -1
else:
a_axis = -1
# TODO(agarwal): handle ndim being None when in graph mode.
b_axis = -2 if b.ndim > 1 else -1
# TODO(srbs): When the shape of the output is a scalar e.g. when performing
# a dot-product of two vectors, numpy returns a scalar object and not an
# instance of ndarray.
# tensordot/MatMul does not support int64 and float64 so we manually cast to
# the compatible types. The conversion may be unsafe.
# TODO(srbs): Figure out why MatMul does not support larger types.
output_type = None
if a.dtype == np.int64:
logging.warning("Unsafe cast to int32.")
a = utils.tensor_to_ndarray(tf.cast(a.data, tf.int32))
b = utils.tensor_to_ndarray(tf.cast(b.data, tf.int32))
output_type = tf.int64
elif a.dtype == np.float64:
logging.warning("Unsafe cast to float32.")
a = utils.tensor_to_ndarray(tf.cast(a.data, tf.float32))
b = utils.tensor_to_ndarray(tf.cast(b.data, tf.float32))
output_type = tf.float64
result_t = tf.tensordot(a.data, b.data, [[a_axis], [b_axis]])
if output_type:
result_t = tf.cast(result_t, output_type)
return utils.tensor_to_ndarray(result_t)
