Backend.arange(..., type_as=...) silently ignores type_as for dtype (and device, in some backends)

[tvercaut]

Backend.arange(stop, start=0, step=1, type_as=None) accepts a type_as argument in every backend, but it is silently ignored (partially or fully) for dtype/device, unlike Backend.ones/zeros/full, which do respect it correctly. This is a footgun: code that reasonably assumes arange follows the same type_as convention as the rest of the backend API (as I did) will silently get the wrong dtype and/or device, with no error or warning.

Repro

import numpy as np
from ot.backend import NumpyBackend

nx = NumpyBackend()
a = np.array([0.1, 0.2, 0.3], dtype=np.float64)

x = nx.arange(3, type_as=a)
print(x.dtype)  # int64 -- expected float64, matching `a`

Compare with nx.ones(3, type_as=a).dtype, which correctly gives float64.

Current behavior per backend (ot/backend.py)

• NumpyBackend.arange (line 1294): return np.arange(start, stop, step) — type_as unused, dtype always inferred by numpy (int for integer bounds).
• JaxBackend.arange (line 1732): return jnp.arange(start, stop, step) — type_as unused.
• TorchBackend.arange (line 2236): sets device=type_as.device but not dtype=type_as.dtype — device is respected, dtype is not.
• CupyBackend.arange (line 2789): return cp.arange(start, stop, step) — type_as unused.
• TensorflowBackend.arange (line 3242): return tnp.arange(start, stop, step) — type_as unused.

None of them pass dtype=type_as.dtype (or device=type_as.device for cupy/tf), unlike e.g. NumpyBackend.ones:

def ones(self, shape, type_as=None):
    if type_as is None:
        return np.ones(shape)
    else:
        return np.ones(shape, dtype=type_as.dtype)

Where I hit this

While prototyping a change to ot/lp/_grid.py (in the still-open #863), I built x = nx.arange(n, type_as=A) to get float grid positions matching histogram A's dtype, then passed it into ot.lp.emd_1d, which derives its output coupling's dtype from the positions array (x_a) rather than the weights. Since x silently stayed int64, the resulting sparse transport plan was silently truncated to all zeros — no exception, no warning, just wrong numbers. I've since worked around it locally by building the position array via nx.cumsum(nx.ones(n, type_as=A)) - 1.0 instead, since ones does respect type_as.

Suggested fix

Make arange respect type_as consistently with ones/zeros/full, e.g. for NumpyBackend:

def arange(self, stop, start=0, step=1, type_as=None):
    if type_as is None:
        return np.arange(start, stop, step)
    else:
        return np.arange(start, stop, step, dtype=type_as.dtype)

and similarly for the other backends (adding device=type_as.device where applicable, and completing TorchBackend.arange's existing partial handling with dtype=type_as.dtype).

[tvercaut]

Implementing 1D gradients in #863 was made easier by fixing this issue. For simplicity, I bundled a fix in this (still open) PR rather than opening a dedicated one. Hope that's OK.