Repository: google/jax

Issue:
cumsum / cumprod int dtype documentation inconsistent with behaviour
### Description

Both `cumsum` and `cumprod` have `dtype` arguments that are documented as:

>dtype : dtype, optional
        Type of the returned array and of the accumulator in which the
        elements are summed.  If `dtype` is not specified, it defaults
        to the dtype of `a`, unless `a` has an integer dtype with a
        precision less than that of the default platform integer.  In
        that case, the default platform integer is used.

However, integer types are not raised when no dtype is specified

```python
import jax.numpy as jnp
print(jnp.cumsum(jnp.asarray([2, 3], dtype=jnp.uint8)).dtype)  # uint8
print(jnp.cumprod(jnp.asarray([2, 3], dtype=jnp.uint8)).dtype)  # uint8
```

Demo on [colab](https://colab.research.google.com/drive/19hJ6am4cRpMPG29IHFKnpmpce1O-eMKx?usp=sharing)

This behaviour (returning the same dtype as the input) makes sense but isn't what is documented. 

### What jax/jaxlib version are you using?

jax 0.4.16, jaxlib 0.4.16

### Which accelerator(s) are you using?

CPU and T4 GPU (colab)

### Additional system info

colab

### NVIDIA GPU info

```
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.105.17   Driver Version: 525.105.17   CUDA Version: 12.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla T4            Off  | 00000000:00:04.0 Off |                    0 |
| N/A   48C    P0    27W /  70W |  11743MiB / 15360MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
+-----------------------------------------------------------------------------+
```


Discussion:
Hi - thanks for the report! It's easy to miss, but the per-parameter documentation does not come from JAX, but from NumPy (look for *original docstring below* at the top of the function documentation).

That said, we do try to note the differences between numpy and JAX in the description, so we should probably add some verbiage about that here.

Resolving patch:
diff --git a/jax/_src/numpy/reductions.py b/jax/_src/numpy/reductions.py
--- a/jax/_src/numpy/reductions.py
+++ b/jax/_src/numpy/reductions.py
@@ -650,9 +650,16 @@ def __call__(self, a: ArrayLike, axis: Axis = None,
                dtype: Optional[DTypeLike] = None, out: None = None) -> Array: ...
 
 
+# TODO(jakevdp): should we change these semantics to match those of numpy?
+CUML_REDUCTION_LAX_DESCRIPTION = """
+Unlike the numpy counterpart, when ``dtype`` is not specified the output dtype will always
+match the dtype of the input.
+"""
+
 def _make_cumulative_reduction(np_reduction: Any, reduction: Callable[..., Array],
                                fill_nan: bool = False, fill_value: ArrayLike = 0) -> CumulativeReduction:
-  @_wraps(np_reduction, skip_params=['out'])
+  @_wraps(np_reduction, skip_params=['out'],
+          lax_description=CUML_REDUCTION_LAX_DESCRIPTION)
   def cumulative_reduction(a: ArrayLike, axis: Axis = None,
                            dtype: Optional[DTypeLike] = None, out: None = None) -> Array:
     return _cumulative_reduction(a, _ensure_optional_axes(axis), dtype, out)


Test patch: