Streamlining JAX's codebase and reducing boilerplate code

[carlosgmartin]

Meta-question inspired by this PR:

Much of JAX's codebase uses boilerplate like check_arraylike (though not always), promote_args, promote_args_inexact, lax.convert_element_type, a jax.jit wrapper around functions (not always), etc.

For examples of these in use, see jax/_src/numpy/ufuncs.py.

Might there be a way to refactor the codebase so that most of this boilerplate can be eliminated?

[jakevdp]

All those helper functions are the result of refactoring the codebase to eliminate boilerplate 😁

But in all seriousness, I'd be happy to revisit these. One thing I'd thought about doing was having a single combined utility to call check_arraylike and convert inputs via jnp.asarray, which is typically done as two steps currently (I proposed this in #14921).

Did you have any specific ideas?

[carlosgmartin]

#14921 sounds great! One example I had in mind from #23923 are the lines

x1, = promote_args_inexact("ldexp", x1)
x2 = lax.convert_element_type(x2, dtypes.dtype(x1))

Would it be possible to collapse these into a single operation? Perhaps something like

x1, x2 = promote_compatible("ldexp", x1, x2)

promote_compatible could apply both shape promotion and dtype promotion to the join of the arguments' types, exact or inexact.

Additionally, I notice that tests use strict dtype promotion:

jax._src.dtypes.TypePromotionError: Input dtypes ('int16', 'int32') have no available implicit dtype promotion path when jax_numpy_dtype_promotion=strict. Try explicitly casting inputs to the desired output type, or set jax_numpy_dtype_promotion=standard.

Is there a reason to use strict promotion, if arrays are going to be promoted manually anyway? (Dropping this would remove the need to explicitly add code for promotion in each function.)

[jakevdp]

That already exists, and you can do it with x1, x2 = promote_args_inexact("ldexp", x1, x2). The reason we didn't use that in #23923 is because we didn't want to use standard promotion rules: it's a function that expects the inputs to be a float and an int, which is not typical of most APIs.

[jakevdp]

Regarding strict promotion: we always use strict promotion in tests, because if our APIs are not tested under strict promotion then the strict promotion would not be useful to users, as it would trigger on internal implementations rather than in user code.

[carlosgmartin]

Could this be done with a wrapper that wraps an API function in a context manager that locally allows standard dtype and rank promotion, just for the implementation of the function itself? Something along the lines of

def use_standard_promotion(f):
  def g(*args, **kwargs):
    with jax.numpy_dtype_promotion("standard"):
      with jax.numpy_rank_promotion("allow"):
        return f(*args, **kwargs)
  return g

@use_standard_promotion
def ldexp(x1, x2):
  ...

[jakevdp]

You can use those context managers directly as decorators for this: @jax.numpy_dtype_promotion('standard').

In the case of #23923 we didn't want to use that, because it's not a case where we want x1 and x2 to promote to a common dtype; rather it's a case where we have an inexact x1 and an integer x2 and we want to cast x2 to the dtype of x1.

Again ldexp is not a typical API, so our typical helper utilities don't apply.