def diagonal(a, offset=0, axis1=0, axis2=1): # pylint: disable=missing-docstring a = asarray(a).data maybe_rank = a.shape.rank if maybe_rank is not None and offset == 0 and ( axis1 == maybe_rank - 2 or axis1 == -2) and (axis2 == maybe_rank - 1 or axis2 == -1): return np_utils.tensor_to_ndarray(array_ops.matrix_diag_part(a)) a = moveaxis(np_utils.tensor_to_ndarray(a), (axis1, axis2), (-2, -1)).data a_shape = array_ops.shape(a) def _zeros(): # pylint: disable=missing-docstring return (array_ops.zeros(array_ops.concat([a_shape[:-1], [0]], 0), dtype=a.dtype), 0) # All zeros since diag_part doesn't handle all possible k (aka offset). # Written this way since cond will run shape inference on both branches, # and diag_part shape inference will fail when offset is out of bounds. a, offset = np_utils.cond( np_utils.logical_or( np_utils.less_equal(offset, -1 * np_utils.getitem(a_shape, -2)), np_utils.greater_equal(offset, np_utils.getitem(a_shape, -1)), ), _zeros, lambda: (a, offset)) a = np_utils.tensor_to_ndarray(array_ops.matrix_diag_part(a, k=offset)) return a
def sort(a, axis=-1, kind='quicksort', order=None): # pylint: disable=missing-docstring if kind != 'quicksort': raise ValueError("Only 'quicksort' is supported.") if order is not None: raise ValueError("'order' argument to sort is not supported.") a = np_array_ops.array(a) if axis is None: result_t = sort_ops.sort(array_ops.reshape(a.data, [-1]), 0) return np_utils.tensor_to_ndarray(result_t) else: return np_utils.tensor_to_ndarray(sort_ops.sort(a.data, axis))
def poisson(lam=1.0, size=None): if size is None: size = () elif np_utils.isscalar(size): size = (size, ) return np_utils.tensor_to_ndarray( random_ops.random_poisson(shape=size, lam=lam, dtype=np_dtypes.int64))
def einsum(subscripts, *operands, **kwargs): # pylint: disable=missing-docstring casting = kwargs.get('casting', 'safe') optimize = kwargs.get('optimize', False) if casting == 'safe': operands = np_array_ops._promote_dtype(*operands) # pylint: disable=protected-access elif casting == 'no': operands = [np_array_ops.asarray(x) for x in operands] else: raise ValueError('casting policy not supported: %s' % casting) if not optimize: # TF doesn't have a "no optimization" option. # TODO(wangpeng): Print a warning that np and tf use different # optimizations. tf_optimize = 'greedy' elif optimize == True: # pylint: disable=singleton-comparison,g-explicit-bool-comparison tf_optimize = 'greedy' elif optimize == 'greedy': tf_optimize = 'greedy' elif optimize == 'optimal': tf_optimize = 'optimal' else: raise ValueError('`optimize` method not supported: %s' % optimize) operands = [x.data for x in operands] res = special_math_ops.einsum(subscripts, *operands, optimize=tf_optimize) res = np_utils.tensor_to_ndarray(res) return res
def _bin_op(tf_fun, a, b, promote=True): if promote: a, b = np_array_ops._promote_dtype(a, b) # pylint: disable=protected-access else: a = np_array_ops.array(a) b = np_array_ops.array(b) return np_utils.tensor_to_ndarray(tf_fun(a.data, b.data))
def diag(v, k=0): # pylint: disable=missing-docstring """Raises an error if input is not 1- or 2-d.""" v = asarray(v).data v_rank = array_ops.rank(v) v.shape.with_rank_at_most(2) # TODO(nareshmodi): Consider a np_utils.Assert version that will fail during # tracing time if the shape is known. control_flow_ops.Assert( np_utils.logical_or(math_ops.equal(v_rank, 1), math_ops.equal(v_rank, 2)), [v_rank]) def _diag(v, k): return np_utils.cond( math_ops.equal(array_ops.size(v), 0), lambda: array_ops.zeros([abs(k), abs(k)], dtype=v.dtype), lambda: array_ops.matrix_diag(v, k=k)) def _diag_part(v, k): v_shape = array_ops.shape(v) v, k = np_utils.cond( np_utils.logical_or( np_utils.less_equal(k, -1 * np_utils.getitem(v_shape, 0)), np_utils.greater_equal(k, np_utils.getitem(v_shape, 1)), ), lambda: (array_ops.zeros([0, 0], dtype=v.dtype), 0), lambda: (v, k)) result = array_ops.matrix_diag_part(v, k=k) return result result = np_utils.cond(math_ops.equal(v_rank, 1), lambda: _diag(v, k), lambda: _diag_part(v, k)) return np_utils.tensor_to_ndarray(result)
def tri(N, M=None, k=0, dtype=None): # pylint: disable=invalid-name,missing-docstring M = M if M is not None else N if dtype is not None: dtype = np_utils.result_type(dtype) else: dtype = np_dtypes.default_float_type() if k < 0: lower = -k - 1 if lower > N: r = array_ops.zeros([N, M], dtype) else: # Keep as tf bool, since we create an upper triangular matrix and invert # it. o = array_ops.ones([N, M], dtype=dtypes.bool) r = math_ops.cast( math_ops.logical_not(array_ops.matrix_band_part(o, lower, -1)), dtype) else: o = array_ops.ones([N, M], dtype) if k > M: r = o else: r = array_ops.matrix_band_part(o, -1, k) return np_utils.tensor_to_ndarray(r)
def split(ary, indices_or_sections, axis=0): ary = asarray(ary) if not isinstance(indices_or_sections, six.integer_types): indices_or_sections = _boundaries_to_sizes(ary, indices_or_sections, axis) result = array_ops.split(ary.data, indices_or_sections, axis=axis) return [np_utils.tensor_to_ndarray(a) for a in result]
def standard_normal(size=None): # TODO(wangpeng): Use new stateful RNG if size is None: size = () elif np_utils.isscalar(size): size = (size,) dtype = np_dtypes.default_float_type() return np_utils.tensor_to_ndarray(random_ops.random_normal(size, dtype=dtype))
def _argminmax(fn, a, axis=None): a = np_array_ops.array(a) if axis is None: # When axis is None numpy flattens the array. a_t = array_ops.reshape(a.data, [-1]) else: a_t = np_array_ops.atleast_1d(a).data return np_utils.tensor_to_ndarray(fn(input=a_t, axis=axis))
def uniform(low=0.0, high=1.0, size=None): dtype = np_dtypes.default_float_type() low = np_array_ops.asarray(low, dtype=dtype) high = np_array_ops.asarray(high, dtype=dtype) if size is None: size = array_ops.broadcast_dynamic_shape(low.shape, high.shape) return np_utils.tensor_to_ndarray( random_ops.random_uniform( shape=size, minval=low, maxval=high, dtype=dtype))
def moveaxis(a, source, destination): # pylint: disable=missing-docstring """Raises ValueError if source, destination not in (-ndim(a), ndim(a)).""" if not source and not destination: return a a = asarray(a).data if isinstance(source, int): source = (source, ) if isinstance(destination, int): destination = (destination, ) a_rank = np_utils._maybe_static(array_ops.rank(a)) # pylint: disable=protected-access def _correct_axis(axis, rank): if axis < 0: return axis + rank return axis source = tuple(_correct_axis(axis, a_rank) for axis in source) destination = tuple(_correct_axis(axis, a_rank) for axis in destination) if a.shape.rank is not None: perm = [i for i in range(a_rank) if i not in source] for dest, src in sorted(zip(destination, source)): assert dest <= len(perm) perm.insert(dest, src) else: r = math_ops.range(a_rank) def _remove_indices(a, b): """Remove indices (`b`) from `a`.""" items = array_ops.unstack(sort_ops.sort(array_ops.stack(b)), num=len(b)) i = 0 result = [] for item in items: result.append(a[i:item]) i = item + 1 result.append(a[i:]) return array_ops.concat(result, 0) minus_sources = _remove_indices(r, source) minus_dest = _remove_indices(r, destination) perm = array_ops.scatter_nd(array_ops.expand_dims(minus_dest, 1), minus_sources, [a_rank]) perm = array_ops.tensor_scatter_update( perm, array_ops.expand_dims(destination, 1), source) a = array_ops.transpose(a, perm) return np_utils.tensor_to_ndarray(a)
def where(condition, x=None, y=None): """Raises ValueError if exactly one of x or y is not None.""" condition = asarray(condition, dtype=np.bool_) if x is None and y is None: return nonzero(condition) elif x is not None and y is not None: x, y = _promote_dtype(x, y) return np_utils.tensor_to_ndarray( array_ops.where_v2(condition.data, x.data, y.data)) raise ValueError('Both x and y must be ndarrays, or both must be None.')
def _comparison(tf_fun, x1, x2, cast_bool_to_int=False): dtype = np_utils.result_type(x1, x2) # Cast x1 and x2 to the result_type if needed. x1 = np_array_ops.array(x1, dtype=dtype) x2 = np_array_ops.array(x2, dtype=dtype) x1 = x1.data x2 = x2.data if cast_bool_to_int and x1.dtype == dtypes.bool: x1 = math_ops.cast(x1, dtypes.int32) x2 = math_ops.cast(x2, dtypes.int32) return np_utils.tensor_to_ndarray(tf_fun(x1, x2))
def clip(a, a_min, a_max): # pylint: disable=missing-docstring if a_min is None and a_max is None: raise ValueError('Not more than one of `a_min` and `a_max` may be `None`.') if a_min is None: return minimum(a, a_max) elif a_max is None: return maximum(a, a_min) else: a, a_min, a_max = np_array_ops._promote_dtype(a, a_min, a_max) # pylint: disable=protected-access return np_utils.tensor_to_ndarray( clip_ops.clip_by_value( *np_utils.tf_broadcast(a.data, a_min.data, a_max.data)))
def expand_dims(a, axis): """Expand the shape of an array. 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`. axis: int. axis on which to expand the shape. Returns: An ndarray with the contents and dtype of `a` and shape expanded on axis. """ a = asarray(a) return np_utils.tensor_to_ndarray(array_ops.expand_dims(a.data, axis=axis))
def cumsum(a, axis=None, dtype=None): # pylint: disable=missing-docstring a = asarray(a, dtype=dtype) if dtype is None: a = _maybe_promote_to_int(a) # If axis is None, the input is flattened. if axis is None: a = ravel(a) axis = 0 elif axis < 0: axis += array_ops.rank(a.data) return np_utils.tensor_to_ndarray(math_ops.cumsum(a.data, axis))
def around(a, decimals=0): # pylint: disable=missing-docstring a = asarray(a) dtype = a.dtype factor = math.pow(10, decimals) # Use float as the working dtype instead of a.dtype, because a.dtype can be # integer and `decimals` can be negative. float_dtype = np_dtypes.default_float_type() a = a.astype(float_dtype).data factor = math_ops.cast(factor, float_dtype) a = math_ops.multiply(a, factor) a = math_ops.round(a) a = math_ops.divide(a, factor) return np_utils.tensor_to_ndarray(a).astype(dtype)
def randint(low, high=None, size=None, dtype=onp.int): # pylint: disable=missing-function-docstring low = int(low) if high is None: high = low low = 0 if size is None: size = () elif isinstance(size, int): size = (size,) dtype = np_utils.result_type(dtype) if dtype not in (onp.int32, onp.int64): raise ValueError('Only np.int32 or np.int64 types are supported') return np_utils.tensor_to_ndarray( random_ops.random_uniform( shape=size, minval=low, maxval=high, dtype=dtype))
def reshape(a, newshape): """Reshapes an array. 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`. newshape: 0-d or 1-d array_like. Returns: An ndarray with the contents and dtype of `a` and shape `newshape`. """ a = asarray(a) if isinstance(newshape, np_arrays.ndarray): newshape = newshape.data return np_utils.tensor_to_ndarray(array_ops.reshape(a.data, newshape))
def swapaxes(a, axis1, axis2): # pylint: disable=missing-docstring a = asarray(a) a_rank = array_ops.rank(a) if axis1 < 0: axis1 += a_rank if axis2 < 0: axis2 += a_rank perm = math_ops.range(a_rank) perm = array_ops.tensor_scatter_update(perm, [[axis1], [axis2]], [axis2, axis1]) a = array_ops.transpose(a, perm) return np_utils.tensor_to_ndarray(a)
def trace(a, offset=0, axis1=0, axis2=1, dtype=None): # pylint: disable=missing-docstring if dtype: dtype = np_utils.result_type(dtype) a = np_array_ops.asarray(a, dtype).data if offset == 0: a_shape = a.shape if a_shape.rank is not None: rank = len(a_shape) if (axis1 == -2 or axis1 == rank - 2) and (axis2 == -1 or axis2 == rank - 1): return np_utils.tensor_to_ndarray(math_ops.trace(a)) a = np_array_ops.diagonal(a, offset, axis1, axis2) return np_array_ops.sum(a, -1, dtype)
def squeeze(a, axis=None): """Removes single-element axes from the array. 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`. axis: scalar or list/tuple of ints. TODO(srbs): tf.squeeze throws error when axis is a Tensor eager execution is enabled. So we cannot allow axis to be array_like here. Fix. Returns: An ndarray. """ a = asarray(a) return np_utils.tensor_to_ndarray(array_ops.squeeze(a, axis))
def imag(a): """Returns imaginary parts of all elements in `a`. Uses `tf.imag`. 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`. Returns: An ndarray with the same shape as `a`. """ a = asarray(a) # TODO(srbs): np.imag returns a scalar if a is a scalar, whereas we always # return an ndarray. return np_utils.tensor_to_ndarray(math_ops.imag(a.data))
def maximum(x1, x2): # pylint: disable=missing-function-docstring # Fast path for when maximum is used as relu. if isinstance(x2, numbers.Real) and not isinstance( x2, bool) and x2 == 0 and isinstance( x1, np_arrays.ndarray) and not x1._is_boolean(): # pylint: disable=protected-access return np_utils.tensor_to_ndarray( nn_ops.relu(np_array_ops.asarray(x1).data)) def max_or_or(x1, x2): if x1.dtype == dtypes.bool: assert x2.dtype == dtypes.bool return math_ops.logical_or(x1, x2) return math_ops.maximum(x1, x2) return _bin_op(max_or_or, x1, x2)
def real(val): """Returns real parts of all elements in `a`. Uses `tf.real`. Args: val: 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 with the same shape as `a`. """ val = asarray(val) # TODO(srbs): np.real returns a scalar if val is a scalar, whereas we always # return an ndarray. return np_utils.tensor_to_ndarray(math_ops.real(val.data))
def randn(*args): """Returns samples from a normal distribution. Uses `tf.random_normal`. Args: *args: The shape of the output array. Returns: An ndarray with shape `args` and dtype `float64`. """ # TODO(wangpeng): Use new stateful RNG if np_utils.isscalar(args): args = (args, ) return np_utils.tensor_to_ndarray( random_ops.random_normal(args, dtype=DEFAULT_RANDN_DTYPE))
def compress(condition, a, axis=None): """Compresses `a` by selecting values along `axis` with `condition` true. Uses `tf.boolean_mask`. Args: condition: 1-d array of bools. If `condition` is shorter than the array axis (or the flattened array if axis is None), it is padded with False. a: array_like. Could be an ndarray, a Tensor or any object that can be converted to a Tensor using `tf.convert_to_tensor`. axis: Optional. Axis along which to select elements. If None, `condition` is applied on flattened array. Returns: An ndarray. Raises: ValueError: if `condition` is not of rank 1. """ condition = asarray(condition, dtype=bool) a = asarray(a) if condition.ndim != 1: raise ValueError('condition must be a 1-d array.') # `np.compress` treats scalars as 1-d arrays. if a.ndim == 0: a = ravel(a) if axis is None: a = ravel(a) axis = 0 if axis < 0: axis += a.ndim assert axis >= 0 and axis < a.ndim # `tf.boolean_mask` requires the first dimensions of array and condition to # match. `np.compress` pads condition with False when it is shorter. condition_t = condition.data a_t = a.data if condition.shape[0] < a.shape[axis]: padding = array_ops.fill([a.shape[axis] - condition.shape[0]], False) condition_t = array_ops.concat([condition_t, padding], axis=0) return np_utils.tensor_to_ndarray( array_ops.boolean_mask(tensor=a_t, mask=condition_t, axis=axis))
def transpose(a, axes=None): """Permutes dimensions of the array. 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`. axes: array_like. A list of ints with length rank(a) or None specifying the order of permutation. The i'th dimension of the output array corresponds to axes[i]'th dimension of the `a`. If None, the axes are reversed. Returns: An ndarray. """ a = asarray(a) if axes is not None: axes = asarray(axes) return np_utils.tensor_to_ndarray(array_ops.transpose(a=a.data, perm=axes))
def ravel(a): """Flattens `a` into a 1-d array. If `a` is already a 1-d ndarray it is returned as is. Uses `tf.reshape`. 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`. Returns: A 1-d ndarray. """ a = asarray(a) if a.ndim == 1: return a return np_utils.tensor_to_ndarray(array_ops.reshape(a.data, [-1]))