jax.dtypes.result_type

jax.dtypes.result_type(*args, return_weak_type_flag=False)

Convenience function to apply JAX argument dtype promotion.

Parameters:

• return_weak_type_flag (bool) – if True, then return a (dtype, weak_type) tuple. If False, just return dtype

• args (Any)

Returns:

dtype or (dtype, weak_type) depending on the value of the return_weak_type argument.

Return type:

DType | tuple[DType, bool]