Improvements in TensorFlow Type Promotion Reduce Bit-Widening Risks

Content Overview

• Effects of turning on the new type promotion
• More consistent and predictable promotion type
• Reduced risk of bit-widening
• tf.Tensor mathematical dunder methods
• tf.Variable in-place ops
• tf.constant implicit conversions
• TF-NumPay Array
• Input Type Inference
• WeakTensor-supporting APIs

Below is a non-exhaustive list of changes that result from turning on the new type promotion.

• More consistent and predictable promotion results.
• Reduced risk of bit-widening.
• tf.Tensor mathematical dunder methods use new type promotion.
• tf.constant can return WeakTensor.
• tf.constant allows implicit conversions when a Tensor input with a dtype different from the dtype arg is passed in.
• tf.Variable in-place ops (assign, assign-add, assign-sub) allow implicit conversions.
• tnp.array(1) and tnp.array(1.0) returns 32-bit WeakTensor.
• WeakTensors will be created and used for WeakTensor-supporting unary and binary API's.

More consistent and predictable promotion results

Using a lattice-based system allows the new type promotion to produce consistent and predictable type promotion results.

Old Type Promotion

Changing the order of operations produces inconsistent results using old type promotion.

# Setup
tnp.experimental_enable_numpy_behavior(dtype_conversion_mode="legacy")
a = np.array(1, dtype=np.int8)
b = tf.constant(1)
c = np.array(1, dtype=np.float16)

# (a + b) + c throws an InvalidArgumentError.
try:
  tf.add(tf.add(a, b), c)
except tf.errors.InvalidArgumentError as e:
  print(f'{type(e)}: {e}')  # InvalidArgumentError

<class 'tensorflow.python.framework.errors_impl.InvalidArgumentError'>: cannot compute AddV2 as input #1(zero-based) was expected to be a int8 tensor but is a int32 tensor [Op:AddV2] name:

# (b + a) + c returns an i32 result.
tf.add(tf.add(b, a), c)  # <tf.Tensor: shape=(), dtype=int32, numpy=3>

<tf.Tensor: shape=(), dtype=int32, numpy=3>

New Type Promotion

New type promotion produces consistent results regardless of the order.

tnp.experimental_enable_numpy_behavior(dtype_conversion_mode="all")
a = np.array(1, dtype=np.int8)
b = tf.constant(1)
c = np.array(1, dtype=np.float16)

WARNING:tensorflow:UserWarning: enabling the new type promotion must happen at the beginning of the program. Please ensure no TF APIs have been used yet.

# (a + b) + c returns a f16 result.
tf.add(tf.add(a, b), c)  # <tf.Tensor: shape=(), dtype=float16, numpy=3.0>

<tf.Tensor: shape=(), dtype=float16, numpy=3.0>

# (b + a) + c also returns a f16 result.
tf.add(tf.add(b, a), c)  # <tf.Tensor: shape=(), dtype=float16, numpy=3.0>

<tf.Tensor: shape=(), dtype=float16, numpy=3.0>

Reduced risk of bit-widening

Old Type Promotion

Old type promotion often resulted in 64-bit results.

tnp.experimental_enable_numpy_behavior(dtype_conversion_mode="legacy")

np.array(3.2, np.float16) + tf.constant(1, tf.int8) + tf.constant(50)  # <tf.Tensor: shape=(), dtype=float64, numpy=54.19921875>

<tf.Tensor: shape=(), dtype=float64, numpy=54.19921875>

New Type Promotion

New type promotion returns results with minimal number of bits necessary.

tnp.experimental_enable_numpy_behavior(dtype_conversion_mode="all")

WARNING:tensorflow:UserWarning: enabling the new type promotion must happen at the beginning of the program. Please ensure no TF APIs have been used yet.

np.array(3.2, np.float16) + tf.constant(1, tf.int8) + tf.constant(50)  # <tf.Tensor: shape=(), dtype=float16, numpy=54.2>

<tf.Tensor: shape=(), dtype=float16, numpy=54.1875>

tf.Tensor mathematical dunder methods

All tf.Tensor mathematical dunder methods will follow the new type promotion.

-tf.constant(5)  # <tf.Tensor: shape=(), dtype=int32, numpy=-5, weak=True>

<tf.Tensor: shape=(), dtype=int32, numpy=-5, weak=True>

tf.constant(5, tf.int16) - tf.constant(1, tf.float32)  # <tf.Tensor: shape=(), dtype=float32, numpy=4.0>

<tf.Tensor: shape=(), dtype=float32, numpy=4.0>

tf.Variable in-place ops

Implicit conversions will be allowed in tf.Variable in-place ops.

tnp.experimental_enable_numpy_behavior(dtype_conversion_mode="all")
a = tf.Variable(10, tf.int32)
a.assign_add(tf.constant(5, tf.int16))  # <tf.Variable shape=() dtype=int32, numpy=15>

WARNING:tensorflow:UserWarning: enabling the new type promotion must happen at the beginning of the program. Please ensure no TF APIs have been used yet.
<tf.Variable 'UnreadVariable' shape=() dtype=int32, numpy=15>

tf.constant implicit conversions

In the old type promotion, tf.constant required an input Tensor to have the same dtype as the dtype argument. However, in the new type promotion, we implicitly convert Tensor to the specified dtype.

tnp.experimental_enable_numpy_behavior(dtype_conversion_mode="all")
a = tf.constant(10, tf.int16)
tf.constant(a, tf.float32)  # <tf.Tensor: shape=(), dtype=float32, numpy=10.0>

WARNING:tensorflow:UserWarning: enabling the new type promotion must happen at the beginning of the program. Please ensure no TF APIs have been used yet.
<tf.Tensor: shape=(), dtype=float32, numpy=10.0>

TF-NumPy Array

tnp.array defaults to i32* and f32* for python inputs using the new type promotion.

tnp.array(1)  # <tf.Tensor: shape=(), dtype=int32, numpy=1, weak=True>

<tf.Tensor: shape=(), dtype=int32, numpy=1, weak=True>

tnp.array(1.0)  # <tf.Tensor: shape=(), dtype=int32, numpy=1, weak=True>

<tf.Tensor: shape=(), dtype=float32, numpy=1.0, weak=True>

Input Type Inference

This is how different inputs' types are inferred in the new type promotion.

• tf.Tensor: Since tf.Tensor has a dtype property, we don't do further inference.
• NumPy types: This includes types like np.array(1), np.int16(1), and np.float. Since NumPy inputs also have a dtype property, we take the dtype property as the result inference type. Note that NumPy defaults to i64 and f64.
• Python scalars/Nested types: This includes types like 1, [1, 2, 3], and (1.0, 2.0).
  • Python int is inferred as i32*.
  • Python float is inferred as f32*.
  • Python complex is inferred as c128*.
• If the input doesn't fall into any of the above categories but has a dtype property, we take the dtype property as the result inference type.

Further Reading

The new type promotion closely resembles JAX-NumPy's type promotion. If you want to know more details about the new type promotion and the design choices, check out the resources below.

• JAX Type Promotion Semantics
• Design of Type Promotion Semantics for JAX
• Old TF-NumPy Promotion Semantics

WeakTensor-supporting APIs

Below is a list of APIs that supports WeakTensor.

For an unary op, this means that if an input with no user-specified type is passed in, it will return a WeakTensor.

For a binary op, it will follow the promotion table here. It may or may not return a WeakTensor depending on the promotion result of the two inputs.

• tf.bitwise.invert
• tf.clip_by_value
• tf.debugging.check_numerics
• tf.expand_dims
• tf.identity
• tf.image.adjust_brightness
• tf.image.adjust_gamma
• tf.image.extract_patches
• tf.image.random_brightness
• tf.image.stateless_random_brightness
• tf.linalg.diag
• tf.linalg.diag_part
• tf.linalg.matmul
• tf.linalg.matrix_transpose
• tf.linalg.tensor_diag_part
• tf.linalg.trace
• tf.math.abs
• tf.math.acos
• tf.math.acosh
• tf.math.add
• tf.math.angle
• tf.math.asin
• tf.math.asinh
• tf.math.atan
• tf.math.atanh
• tf.math.ceil
• tf.math.conj
• tf.math.cos
• tf.math.cosh
• tf.math.digamma
• tf.math.divide_no_nan
• tf.math.divide
• tf.math.erf
• tf.math.erfc
• tf.math.erfcinv
• tf.math.erfinv
• tf.math.exp
• tf.math.expm1
• tf.math.floor
• tf.math.floordiv
• tf.math.floormod
• tf.math.imag
• tf.math.lgamma
• tf.math.log1p
• tf.math.log_sigmoid
• tf.math.log
• tf.math.multiply_no_nan
• tf.math.multiply
• tf.math.ndtri
• tf.math.negative
• tf.math.pow
• tf.math.real
• tf.math.real
• tf.math.reciprocal_no_nan
• tf.math.reciprocal
• tf.math.reduce_euclidean_norm
• tf.math.reduce_logsumexp
• tf.math.reduce_max
• tf.math.reduce_mean
• tf.math.reduce_min
• tf.math.reduce_prod
• tf.math.reduce_std
• tf.math.reduce_sum
• tf.math.reduce_variance
• tf.math.rint
• tf.math.round
• tf.math.rsqrt
• tf.math.scalar_mul
• tf.math.sigmoid
• tf.math.sign
• tf.math.sin
• tf.math.sinh
• tf.math.softplus
• tf.math.special.bessel_i0
• tf.math.special.bessel_i0e
• tf.math.special.bessel_i1
• tf.math.special.bessel_i1e
• tf.math.special.bessel_j0
• tf.math.special.bessel_j1
• tf.math.special.bessel_k0
• tf.math.special.bessel_k0e
• tf.math.special.bessel_k1
• tf.math.special.bessel_k1e
• tf.math.special.bessel_y0
• tf.math.special.bessel_y1
• tf.math.special.dawsn
• tf.math.special.expint
• tf.math.special.fresnel_cos
• tf.math.special.fresnel_sin
• tf.math.special.spence
• tf.math.sqrt
• tf.math.square
• tf.math.subtract
• tf.math.tan
• tf.math.tanh
• tf.nn.depth_to_space
• tf.nn.elu
• tf.nn.gelu
• tf.nn.leaky_relu
• tf.nn.log_softmax
• tf.nn.relu6
• tf.nn.relu
• tf.nn.selu
• tf.nn.softsign
• tf.nn.space_to_depth
• tf.nn.swish
• tf.ones_like
• tf.realdiv
• tf.reshape
• tf.squeeze
• tf.stop_gradient
• tf.transpose
• tf.truncatediv
• tf.truncatemod
• tf.zeros_like
• tf.experimental.numpy.abs
• tf.experimental.numpy.absolute
• tf.experimental.numpy.amax
• tf.experimental.numpy.amin
• tf.experimental.numpy.angle
• tf.experimental.numpy.arange
• tf.experimental.numpy.arccos
• tf.experimental.numpy.arccosh
• tf.experimental.numpy.arcsin
• tf.experimental.numpy.arcsinh
• tf.experimental.numpy.arctan
• tf.experimental.numpy.arctanh
• tf.experimental.numpy.around
• tf.experimental.numpy.array
• tf.experimental.numpy.asanyarray
• tf.experimental.numpy.asarray
• tf.experimental.numpy.ascontiguousarray
• tf.experimental.numpy.average
• tf.experimental.numpy.bitwise_not
• tf.experimental.numpy.cbrt
• tf.experimental.numpy.ceil
• tf.experimental.numpy.conj
• tf.experimental.numpy.conjugate
• tf.experimental.numpy.copy
• tf.experimental.numpy.cos
• tf.experimental.numpy.cosh
• tf.experimental.numpy.cumprod
• tf.experimental.numpy.cumsum
• tf.experimental.numpy.deg2rad
• tf.experimental.numpy.diag
• tf.experimental.numpy.diagflat
• tf.experimental.numpy.diagonal
• tf.experimental.numpy.diff
• tf.experimental.numpy.empty_like
• tf.experimental.numpy.exp2
• tf.experimental.numpy.exp
• tf.experimental.numpy.expand_dims
• tf.experimental.numpy.expm1
• tf.experimental.numpy.fabs
• tf.experimental.numpy.fix
• tf.experimental.numpy.flatten
• tf.experimental.numpy.flip
• tf.experimental.numpy.fliplr
• tf.experimental.numpy.flipud
• tf.experimental.numpy.floor
• tf.experimental.numpy.full_like
• tf.experimental.numpy.imag
• tf.experimental.numpy.log10
• tf.experimental.numpy.log1p
• tf.experimental.numpy.log2
• tf.experimental.numpy.log
• tf.experimental.numpy.max
• tf.experimental.numpy.mean
• tf.experimental.numpy.min
• tf.experimental.numpy.moveaxis
• tf.experimental.numpy.nanmean
• tf.experimental.numpy.negative
• tf.experimental.numpy.ones_like
• tf.experimental.numpy.positive
• tf.experimental.numpy.prod
• tf.experimental.numpy.rad2deg
• tf.experimental.numpy.ravel
• tf.experimental.numpy.real
• tf.experimental.numpy.reciprocal
• tf.experimental.numpy.repeat
• tf.experimental.numpy.reshape
• tf.experimental.numpy.rot90
• tf.experimental.numpy.round
• tf.experimental.numpy.signbit
• tf.experimental.numpy.sin
• tf.experimental.numpy.sinc
• tf.experimental.numpy.sinh
• tf.experimental.numpy.sort
• tf.experimental.numpy.sqrt
• tf.experimental.numpy.square
• tf.experimental.numpy.squeeze
• tf.experimental.numpy.std
• tf.experimental.numpy.sum
• tf.experimental.numpy.swapaxes
• tf.experimental.numpy.tan
• tf.experimental.numpy.tanh
• tf.experimental.numpy.trace
• tf.experimental.numpy.transpose
• tf.experimental.numpy.triu
• tf.experimental.numpy.vander
• tf.experimental.numpy.var
• tf.experimental.numpy.zeros_like