Type promotion semantics#

This document describes JAX’s type promotion rules–i.e., the result ofjax.numpy.promote_types() for each pair of types.For some background on the considerations that went into the design of what is described below, seeDesign of Type Promotion Semantics for JAX.

JAX’s type promotion behavior is determined via the following type promotion lattice:

where, for example:

• b1 meansnp.bool_,

• i2 meansnp.int16,

• u4 meansnp.uint32,

• bf meansnp.bfloat16,

• f2 meansnp.float16,

• c8 meansnp.complex64,

• i* means Pythonint or weakly-typedint,

• f* means Pythonfloat or weakly-typedfloat, and

• c* means Pythoncomplex or weakly-typedcomplex.

(for more about weak types, seeWeakly-typed values in JAX below).

Promotion between any two types is given by theirjoinon this lattice, which generates the following binary promotion table:

Jax’s type promotion rules differ from those of NumPy, as given bynumpy.promote_types(), in those cells highlighted with a green backgroundin the table above. There are three key classes of differences:

• When promoting a weakly typed value against a typed JAX value of the same category,JAX always prefers the precision of the JAX value. For example,jnp.int16(1)+1will returnint16 rather than promoting toint64 as in NumPy.Note that this applies only to Python scalar values; if the constant is a NumPyarray then the above lattice is used for type promotion.For example,jnp.int16(1)+np.array(1) will returnint64.

• When promoting an integer or boolean type against a floating-point or complextype, JAX always prefers the type of the floating-point or complex type.

• JAX supports thebfloat16non-standard 16-bit floating point type(jax.numpy.bfloat16), which is useful for neural network training.The only notable promotion behavior is with respect to IEEE-754float16, with whichbfloat16 promotes to afloat32.

The differences between NumPy and JAX are motivated by the fact thataccelerator devices, such as GPUs and TPUs, either pay a significantperformance penalty to use 64-bit floating point types (GPUs) or do notsupport 64-bit floating point types at all (TPUs). Classic NumPy’s promotionrules are too willing to overpromote to 64-bit types, which is problematic fora system designed to run on accelerators.

JAX uses floating point promotion rules that are more suited to modernaccelerator devices and are less aggressive about promoting floating pointtypes. The promotion rules used by JAX for floating-point types are similar tothose used by PyTorch.

>>>x=jnp.arange(5,dtype='int8')>>>2*xArray([0, 2, 4, 6, 8], dtype=int8)

>>>jnp.int32(2)*xArray([0, 2, 4, 6, 8], dtype=int32)

>>>jnp.asarray(2)Array(2, dtype=int32, weak_type=True)

>>>jnp.asarray(2,dtype='int32')Array(2, dtype=int32)

>>>x=jnp.float32(1)>>>y=jnp.int32(1)>>>withjax.numpy_dtype_promotion('strict'):...z=x+y...Traceback (most recent call last):TypePromotionError:Input dtypes ('float32', 'int32') have no available implicitdtype promotion path when jax_numpy_dtype_promotion=strict. Try explicitly castinginputs to the desired output type, or set jax_numpy_dtype_promotion=standard.

>>>withjax.numpy_dtype_promotion('strict'):...z=x+1>>>print(z)2.0

jax.config.update('jax_numpy_dtype_promotion','strict')

jax.config.update('jax_numpy_dtype_promotion','standard')