Repository: google/jax
Issue: jnp.promote_dtype does not preserve weak types

Assuming the type promotion lattice presented in the [great documentation](https://jax.readthedocs.io/en/latest/type_promotion.html) by @jakevdp, I'd assume that the promotion of two python-native `float*` would be a new `float*`. 
This is confirmed by the table  in the same page.

However, 
```
>>> jnp.promote_types(float, float)
dtype('float64')                                   # I expected float*
```

I do understand that there might be internal reasons why this is useful. I guess that internally jax uses this API to do some type promotion before calling lax primitives.

However, if someone tries to use this API to build objects that respect the same type-promotion semantics as Jax's arrays (like I need to do in NetKet for some lazy objects representing sums of matrices), I cannot use `jnp.promote_types` because it discards the `weak` attribute and promotes everything to non-weak types.
 
I acknowledge it's not a mainstream use-case, but it would be useful if Jax exposed an API that performs the same type promotion as jax does internally. 

Changed files:
jax/_src/dtypes.py, tests/dtypes_test.py

diff --git a/jax/_src/dtypes.py b/jax/_src/dtypes.py
index a8caf1c42..96bb22ffd 100644
--- a/jax/_src/dtypes.py
+++ b/jax/_src/dtypes.py
@@ -375,11 +375,20 @@ def _lattice_result_type(*args):
     result_type = _least_upper_bound(*{_jax_type(d, w) for d, w in zip(dtypes, weak_types)})
     return dtype(result_type), any(result_type is t for t in _weak_types)
 
-def result_type(*args):
-  """Convenience function to apply JAX argument dtype promotion."""
+def result_type(*args, return_weak_type_flag=False):
+  """Convenience function to apply JAX argument dtype promotion.
+
+  Args:
+    return_weak_type_flag : if True, then return a ``(dtype, weak_type)`` tuple.
+      If False, just return `dtype`
+
+  Returns:
+    dtype or (dtype, weak_type) depending on the value of the ``return_weak_type`` argument.
+  """
   if len(args) == 0:
     raise ValueError("at least one array or dtype is required")
   dtype, weak_type = _lattice_result_type(*(float_ if arg is None else arg for arg in args))
   if weak_type:
     dtype = _default_types['f' if dtype == _bfloat16_dtype else dtype.kind]
-  return canonicalize_dtype(dtype)
+  dtype = canonicalize_dtype(dtype)
+  return (dtype, weak_type) if return_weak_type_flag else dtype
diff --git a/tests/dtypes_test.py b/tests/dtypes_test.py
index ec29617aa..231cfdac0 100644
--- a/tests/dtypes_test.py
+++ b/tests/dtypes_test.py
@@ -265,6 +265,15 @@ class TestPromotionTables(jtu.JaxTestCase):
     # This matches the behavior of np.result_type(None) => np.float_
     self.assertEqual(dtypes.result_type(None), dtypes.canonicalize_dtype(dtypes.float_))
 
+  def testResultTypeWeakFlag(self):
+    float_ = dtypes.canonicalize_dtype(dtypes.float_)
+    x_weak = jnp.array(1.)
+    x_strong = x_weak.astype(float_)
+    self.assertEqual(dtypes.result_type(x_weak), float_)
+    self.assertEqual(dtypes.result_type(x_weak, return_weak_type_flag=True), (float_, True))
+    self.assertEqual(dtypes.result_type(x_strong), float_)
+    self.assertEqual(dtypes.result_type(x_strong, return_weak_type_flag=True), (float_, False))
+
   @jtu.ignore_warning(category=UserWarning,
                       message="Explicitly requested dtype.*")
   def testObservedPromotionTable(self):