def _canonicalize_jit_arg(x):
if isinstance(x, ndarray):
return x.data
else:
try:
# We need to convert `int` to the most precise dtype, otherwise the dtype
# of the result may be different from numpy's. For example, when a binary
# op takes in a Python integer 5 and an array of uint32, numpy will pick
# uint32 as 5's dtype, while tf.convert_to_tensor will choose int32 which
# will cause the two arguments to be promoted to int64. We pick uint8
# here, which will be promoted to uint32 by the binary op.
# Note that we prefer unsigned int to signed int when both are equally
# precise. For example, for 5, we pick uint8 instead of int8. There is no
# reason to prefer one to the other, because for each there is a case
# where the behavior diverges from numpy. If we prefer signed int,
# consider the case where the first operand is 5 and the second is
# 2**64-1. Numpy picks uint64 as the result dtype, but because we choose a
# signed type for 5 such as int8, the result type will be float64. On the
# other hand, if we prefer unsigned int, consider the case where the first
# operand is 2**31-1 and the second is -1. Numpy will pick int32, but
# because we choose uint32 for 2*32-1, the result will be int64. The root
# of the problem is that `jit` converts `int` to tensors (hence committing
# to a dtype) too early, when we don't have enough information about the
# jitted function (e.g. which subset of the arguments should be promoted
# together using np.result_type). tf.function doesn't have this problem
# because it doesn't convert `int` to tensors. jax.jit doesn't have this
# problem because it converts `int` to "int tracer" which doesn't commit
# to a dtype.
# TODO(wangpeng): Revisit this design and see whether we can improve `jit`
# and tf.function.
dtype = most_precise_int_dtype(x)
return arrays.convert_to_tensor(value=x, dtype=dtype)
except (TypeError, ValueError):
return x
def array(val, dtype=None, copy=True, ndmin=0): # pylint: disable=redefined-outer-name
"""Creates an ndarray with the contents of val.
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`.
dtype: Optional, defaults to dtype of the `val`. The type of the resulting
ndarray. Could be a python type, a NumPy type or a TensorFlow `DType`.
copy: Determines whether to create a copy of the backing buffer. Since
Tensors are immutable, a copy is made only if val is placed on a different
device than the current one. Even if `copy` is False, a new Tensor may
need to be built to satisfy `dtype` and `ndim`. This is used only if `val`
is an ndarray or a Tensor.
ndmin: The minimum rank of the returned array.
Returns:
An ndarray.
"""
if dtype:
dtype = utils.result_type(dtype)
if isinstance(val, arrays_lib.ndarray):
result_t = val.data
else:
result_t = val
if copy and isinstance(result_t, tf.Tensor):
# Note: In eager mode, a copy of `result_t` is made only if it is not on
# the context device.
result_t = tf.identity(result_t)
if not isinstance(result_t, tf.Tensor):
if not dtype:
dtype = utils.result_type(result_t)
# We can't call `convert_to_tensor(result_t, dtype=dtype)` here because
# convert_to_tensor doesn't allow incompatible arguments such as (5.5, int)
# while np.array allows them. We need to convert-then-cast.
def maybe_data(x):
if isinstance(x, arrays_lib.ndarray):
return x.data
return x
# Handles lists of ndarrays
result_t = tf.nest.map_structure(maybe_data, result_t)
result_t = arrays_lib.convert_to_tensor(result_t)
result_t = tf.cast(result_t, dtype=dtype)
elif dtype:
result_t = tf.cast(result_t, dtype)
ndims = tf.rank(result_t)
def true_fn():
old_shape = tf.shape(result_t)
new_shape = tf.concat([tf.ones(ndmin - ndims, tf.int32), old_shape],
axis=0)
return tf.reshape(result_t, new_shape)
result_t = utils.cond(utils.greater(ndmin, ndims), true_fn,
lambda: result_t)
return arrays_lib.tensor_to_ndarray(result_t)
def bernoulli(key, mean=np.float32(0.5), shape=()):
"""Sample Bernoulli random values with given shape and mean.
Args:
key: a random key, not used in the TF backend (stored in graph).
mean: optional, an array_like broadcastable to `shape` for the mean of the
random variables (default 0.5).
shape: optional, a tuple of nonnegative integers representing the shape
(default scalar).
Returns:
A random array with the specified shape and boolean dtype.
"""
# TODO(wangpeng): convert types TF <-> numpy.
shape = shape or arrays.convert_to_tensor(value=mean).shape
return array(tf.less(uniform(key, shape), mean), copy=False)
def array(val, dtype=None, copy=True, ndmin=0):
"""Creates an ndarray with the contents of val.
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`.
dtype: Optional, defaults to dtype of the `val`. The type of the
resulting ndarray. Could be a python type, a NumPy type or a TensorFlow
`DType`.
copy: Determines whether to create a copy of the backing buffer. Since
Tensors are immutable, a copy is made only if val is placed on a different
device than the current one. Even if `copy` is False, a new Tensor may
need to be built to satisfy `dtype` and `ndim`. This is used only if `val`
is an ndarray or a Tensor.
ndmin: The minimum rank of the returned array.
Returns:
An ndarray.
"""
if dtype:
dtype = utils.result_type(dtype)
if isinstance(val, arrays.ndarray):
result_t = val.data
else:
result_t = val
if copy and isinstance(result_t, tf.Tensor):
# Note: In eager mode, a copy of `result_t` is made only if it is not on
# the context device.
result_t = tf.identity(result_t)
if not isinstance(result_t, tf.Tensor):
if not dtype:
dtype = utils.result_type(result_t)
# We can't call `convert_to_tensor(result_t, dtype=dtype)` here because
# convert_to_tensor doesn't allow incompatible arguments such as (5.5, int)
# while np.array allows them. We need to convert-then-cast.
result_t = arrays.convert_to_tensor(result_t)
result_t = tf.cast(result_t, dtype=dtype)
elif dtype:
result_t = tf.cast(result_t, dtype)
ndims = len(result_t.shape)
if ndmin > ndims:
old_shape = list(result_t.shape)
new_shape = [1 for _ in range(ndmin - ndims)] + old_shape
result_t = tf.reshape(result_t, new_shape)
return arrays.tensor_to_ndarray(result_t)
