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. 


Discussion:
Thanks for the question - this is due to an implementation detail. `jnp.promote_types` is JAX's implementation of numpy's `np.promote_types`, which returns a dtype. A dtype cannot be weak – `weak_type` is a property of JAX arrays, not of dtypes, and so there's no way for `jnp.promote_types` to return a weak type. It's not allowed by the API, which is defined by numpy, not by JAX.

Nevertheless, in binary operations between weak types, you'll see that weakness is preserved::
```python
>>> jnp.add(1.0, 1.0)
DeviceArray(2., dtype=float32, weak_type=True)
```
Ah, I see.
That's unfortunate.

I'm essentially trying to implement something like:

```python
@pytree
class MatrixSum:
    def __init__(self, mat1: ArrayLike, mat2: ArrayLike, dtype=None  ):
         if dtype is None:
             dtype = jnp.promote_dtypes(mat1, mat2)
          self._dtype = dtype
          self._mat1 = jnp.asarray(mat1, dtype=dtype)
          self._mat2 = jnp.asarray(mat2, dtype=dtype)

    @property
    def dtype(self):
         return self._dtype

    def compute_something(self) -> Array[self.dtype]:
        # return a Jax array that might be weakly typed if mat1 and mat2 where.
```

And I wish this behaved (at least within the boundaries of my API) as a jax array, propagating weak and strong dtypes of the input `mat1` and `mat2`. 
But I see that it's impossible to do that within the realms of only using `dtype`. I'll have to find a new solution...

There are APIs within `jax.dtypes` that return `dtype, weak_type` pairs for type promotion. That's essentially what JAX uses to implement binary type promotion.
thanks...

I think what I'd need would be to use `jax._src.dtypes._lattice_result_type` (which is not exported in `jax.dtypes`) and to be able to declare my types as weakly typed, for example by adding a property `MyType.weak_type` to an object

Right now, if an object has a `dtype` property, numpy correctly recognises it. That's not the same for jax.

```python
>>> class Test:
...   def __init__(self, dtype, weak):
...     self.dtype = dtype
...     self.weak_type = weak
...
>>> t = Test(np.float32, True)
>>> np.result_type(t, np.float64)
dtype('float64')
>>> jnp.result_type(t, np.float64)
dtype('float64')
>>> jax._src.dtypes.is_weakly_typed(t)
False
```

you can trick jax by adding an `MyType.aval.weak_type` field but that's a bit convoluted

```python
>>> t.aval = t
>>> jax._src.dtypes.is_weakly_typed(t)
True
```
In short, what I am saying is that it would be useful if jax stabilised the `weak_dtype` interface by

 - recognising a field `.weak_type` in objects much like it recognises `.dtype` to the extent of type promotion
 - exposeed `jax._src.dtypes._lattice_result_type` as a public API, or a similar function that returns both the resulting dtype and the `weak_dtype` boolean flag.
Thanks, I agree we should have some pulic API for dtype promotion that preserves weak types. What would you think about adding  keyword to `jax.dtypes.promote_types` that would allow it to return a `(dtype, weak_type)` pair?

Regarding the second request ("recognizing a field `.weak_type` in objects") can you say more? I'm not sure I understand exactly what you have in mind. What functions should recognize this attribute, and in what context?
>Thanks, I agree we should have some pulic API for dtype promotion that preserves weak types. What would you think about adding keyword to `jax.dtypes.promote_types` that would allow it to return a (dtype, weak_type) pair?

I would very much like that, yes. 

> Regarding the second request ("recognizing a field .weak_type in objects") can you say more? I'm not sure I understand exactly what you have in mind. What functions should recognize this attribute, and in what context?

Right now jax and numpy `(j)np.result_type` works with arbitrary types, as long as they expose a `CustomType.dtype` attribute (see the example above, with the custom type `Test`.
I would propose that `jax.dtypes.promote_types` should define a similar protocol, where it will check for the attribute `CustomType.weak_type` to determine if an object has a weak type, as well as the `.dtype` attribute that it already checks as part of numpy api.

Note that what I propose would be necessary to make `jax.dtypes.promote_types` extensible and compatible with user-defined types. 

For example:
```python
>>> class Test:
...   def __init__(self, dtype, weak):
...     self.dtype = dtype
...     self.weak_type = weak
...
>>> t = Test(np.float64, True, return_weak_flag=True)
>>> jax.dtypes.promote_types(t, jnp.array(3.0, dtype=jnp.float32))
(jnp.float32, False)
>>> jax.dtypes.promote_types(t, jnp.array(3.0), return_weak_flag=True)
(jnp.float64, True)
```

Hmm, at first glance, I'm not sure I like the idea of cementing a `weak_type` attribute into the public JAX API. While it's true that jax arrays have a `weak_type` attribute at the moment, that is an implementation detail rather than something that is intended to be a load-bearing API. For example, we've talked about moving JAX to its own dtype system where weak dtypes are actual dtypes, and the `weak_type` attribute can go away.

A few other options that might work and be more future-proof:

- we could use something like the existing `__jax_array__` method within type promotion of custom objects. Then you could have it return an appropriate jax array
- we could let custom objects export an `aval` attribute that will return a jax abstract value, and pre-empt the `dtype` attribute
- we could provide some sort of hook that would allow `dtypes.is_weakly_typed` to work correctly with custom objects; perhaps delegating to a method like `_is_weakly_typed`

What do you think?
>we could provide some sort of hook that would allow dtypes.is_weakly_typed to work correctly with custom objects; perhaps delegating to a method like _is_weakly_typed

I think that's the best option.
You can add some API that can be used to customise how `_is_weakly_typed` behaves with custom types for now. 
Eventually you can even deprecate it  if you move to tax-specific dtypes (which I agree would be the most elegant solution).

One more comment here: in my opinion this is a misuse of the `promote_dtypes` API. The docstring specifies that the inputs should be a "dtype or dtype specifier". The fact that this happens to accept arbitrary objects that have a `dtype` attribute is unintended, I think, and you probably shouldn't rely on that. Perhaps `result_type` would be a better API to use in this case, and it would make more sense I think to add a `return_weak_flag` there, because the inputs can be arbitrary objects rather than more narrow "dtype or dtype specifier" (and it already calls `is_weakly_typed` on the inputs, so that would be a better place to hook-in any customization of this behavior).
I do agree with you. 
Since Dtypes cannot have Weak dtype attached at the moment, I think that whatever API Jax exposes now should accept objects, not dtypes.
Just to be explicit, since I'm not sure from your reply if you saw it: JAX does already expose `jnp.result_type`, which is, I think, the correct API for type promotion of arbitrary objects, and calls `dtypes.is_weakly_typed` on those objects. We'd just need to add a way for custom objects to tell `is_weakly_typed` what to return.

Resolving patch:
diff --git a/jax/_src/dtypes.py b/jax/_src/dtypes.py
--- 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


Test patch:
diff --git a/tests/dtypes_test.py b/tests/dtypes_test.py
--- a/tests/dtypes_test.py
+++ b/tests/dtypes_test.py
@@ -265,6 +265,15 @@ def testResultTypeNone(self):
     # 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):