# How to use the trax.tf_numpy.numpy.utils.tensor_to_ndarray function in trax

## To help you get started, we’ve selected a few trax examples, based on popular ways it is used in public projects.

Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.

google / trax / trax / tf_numpy / numpy / array_methods.py View on Github
``````# Algorithm:
# 1. Calculate cumulative sum of repeats.
repeats_cumsum = cumsum(repeats)  # [3, 4, 6]
# 2. Use `scatter_nd` to generate an indices list for use in `tf.gather`.
scatter_indices_t = repeats_cumsum[:-1].data  # [3, 4]
scatter_indices_t = tf.expand_dims(scatter_indices_t, 1)  # [[3], [4]]
scatter_updates_t = tf.ones([len(repeats) - 1], dtype=tf.int32)  # [1, 1]
scatter_shape_t = ravel(repeats_cumsum[-1]).data  # [6]
#    `tf.scatter_nd([[3], [4]], [1, 1], [6])` -> `[0, 0, 0, 1, 1, 0]`
scatter_shape_t)
indices_t = tf.cumsum(indices_t)  # [0, 0, 0, 1, 2, 2]
# 3. Use `tf.gather` to gather indices along `axis`.
result_t = tf.gather(a, indices_t, axis=axis)

return utils.tensor_to_ndarray(result_t)``````
google / trax / trax / tf_numpy / numpy / array_creation.py View on Github
``````to power of `base`.
num: Number of values to sample. Defaults to 50.
endpoint: When to include `base ** stop` in the output. Defaults to true.
base: Base of the log space.
dtype: Optional. Type of the resulting ndarray. Could be a python type, a
NumPy type or a TensorFlow `DType`. If not provided, it is figured from
input args.
"""
# TODO(srbs): Check whether dtype is handled properly.
if dtype:
dtype = utils.to_tf_type(dtype)
result = linspace(start, stop, num=num, endpoint=endpoint)
result = tf.pow(base, result.data)
if dtype:
result = utils.maybe_cast(result, dtype)
return utils.tensor_to_ndarray(result)
``````
google / trax / trax / tf_numpy / numpy / math.py View on Github
``````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 = 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.')``````
google / trax / trax / tf_numpy / numpy / array_methods.py View on Github
``````"""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 = array_creation.asarray(a)
if axes is not None:
axes = array_creation.asarray(axes)
return utils.tensor_to_ndarray(tf.transpose(a=a.data, perm=axes))``````
google / trax / trax / tf_numpy / numpy / array_creation.py View on Github
``````raise ValueError('stop: {} must be non-zero.'.format(stop))
if np_sign(start) != np_sign(stop):
raise ValueError('start: {} and stop: {} must have same sign.'.format(
start, stop))
step = 1.
if endpoint:
if num > 1:
step = tf.pow((stop / start), 1 / (num - 1))
else:
step = tf.pow((stop / start), 1 / num)
result = tf.cast(tf.range(num), step.dtype)
result = tf.pow(step, result)
result = tf.multiply(result, start)
if dtype:
result = tf.cast(result, dtype=dtype)
return utils.tensor_to_ndarray(result)``````
google / trax / trax / tf_numpy / numpy / array_methods.py View on Github
``````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 = array_creation.asarray(a)
return utils.tensor_to_ndarray(tf.squeeze(a, axis))``````
google / trax / trax / tf_numpy / numpy / logic.py View on Github
``````to the other.

Args:
x1: array_like. Could be an ndarray, a Tensor or any object that can
be converted to a Tensor using `tf.convert_to_tensor`.
x2: 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 of type bool and broadcasted shape of x1 and x2.
"""
dtype = utils.result_type(x1, x2)
# Cast x1 and x2 to the result_type if needed.
x1 = array_creation.array(x1, copy=False, dtype=dtype)
x2 = array_creation.array(x2, copy=False, dtype=dtype)
return utils.tensor_to_ndarray(tf.equal(x1.data, x2.data))
``````
google / trax / trax / tf_numpy / numpy / array_methods.py View on Github
``````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 = array_creation.asarray(a)
if a.ndim == 1:
return a
return utils.tensor_to_ndarray(tf.reshape(a.data, [-1]))
``````
google / trax / trax / tf_numpy / numpy / math.py View on Github
``````"""Computes the tf_fn(x) for each element in `x`.

Args:
x: array_like. Could be an ndarray, a Tensor or any object that can
be converted to a Tensor using `tf.convert_to_tensor`.
tf_fn: function that takes a single Tensor argument.

Returns:
An ndarray with the same shape as `x`. The default output dtype is
determined by `dtypes.default_float_type`, unless x is an ndarray with a
floating point type, in which case the output type is same as x.dtype.
"""
x = array_creation.asarray(x)
if x.dtype not in (np.float16, np.float32, np.float64):
x = x.astype(dtypes.default_float_type())
return utils.tensor_to_ndarray(tf_fn(x.data))``````
google / trax / trax / tf_numpy / numpy / math.py View on Github
``````"""Computes the max of all array elements or along specified axes.

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: Optional 0-d or 1-d array_like. Axes along which to compute the max.
If None, returns the max of all elements in array.
keepdims: If true, retains reduced dimensions with length 1.

Returns:
An ndarray with the same dtype as `a`.
"""
# TODO(wangpeng): check that we fully match numpy behavior.
a = array_creation.asarray(a)

return utils.tensor_to_ndarray(tf.reduce_max(input_tensor=a.data, axis=axis,
keepdims=keepdims))``````

Trax

Apache-2.0