jax

42dd736a - Change scalar promotion rules to prefer array types over scalar types. (#1709)

Change scalar promotion rules to prefer array types over scalar types. (#1709) * Change scalar promotion rules to prefer array types over scalar types. Currently JAX does not treat Python scalars specially during type promotion. This means that, for example: `1. + np.array([...], np.float32)` ends up as an array of type np.float64. The `1.` is promoted to a default type (here np.float64), and the type promotion of a np.float64 and an np.float32 is an np.float64. This is unlike classic NumPy, which treats scalars specially during type promotion, in particular, preferring the type of an array over the type of a scalar. This change adds a notion of weak_type to JAX avals. During type promotion, we prefer non-weak types, i.e., the type of the array in the example above, ignoring the type of the scalar. In contexts where a Python scalar is to be promoted to a NumPy value, a default type is used (e.g., `np.float_`). This change also makes it possible to use 32-bit default types that differ from NumPy's default types. The JAX test suite passes with 32-bit default types. However, we do not yet enable this change or expose it in the API.