Rank promotion warning#
NumPy broadcasting rules allow the automatic promotion of arguments from one rank (number of array axes) to another. This behavior can be convenient when intended but can also lead to surprising bugs where a silent rank promotion masks an underlying shape error.
Here’s an example of rank promotion:
>>> import numpy as np >>> x = np.arange(12).reshape(4, 3) >>> y = np.array([0, 1, 0]) >>> x + y array([[ 0, 2, 2], [ 3, 5, 5], [ 6, 8, 8], [ 9, 11, 11]])
To avoid potential surprises,
jax.numpy is configurable so that
expressions requiring rank promotion can lead to a warning, error, or can be
allowed just like regular NumPy. The configuration option is named
jax_numpy_rank_promotion and it can take on string values
raise. The default setting is
allow, which allows rank promotion without warning or error.
raise setting raises an error on rank promotion, and
raises a warning on the first occurrence of rank promotion.
Rank promotion can be enabled or disabled locally with the
with jax.numpy_rank_promotion("warn"): z = x + y
This configuration can also be set globally in several ways.
One is by using
jax.config in your code:
from jax import config config.update("jax_numpy_rank_promotion", "warn")
You can also set the option using the environment variable
JAX_NUMPY_RANK_PROMOTION, for example as
JAX_NUMPY_RANK_PROMOTION='warn'. Finally, when using
the option can be set with a command-line flag.