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)