Implement torch.linspace and torch.logspace

experimental/torch_xla2/torch_xla2/ops/jtorch.py

@@ -15,6 +15,16 @@
 from torch_xla2.ops import op_base, mappings, jaten
 import torch_xla2.tensor
 
+TORCH_TO_JAX_DTYPE = {
+    torch.bool: jnp.bool_,
+    torch.float32: jnp.float32,
+    torch.float64: jnp.float64,
+    torch.int64: jnp.int64,
+    torch.int32: jnp.int32,
+    torch.complex64: jnp.complex64,
+    # add more as needed
+}
+
 
 def register_function(torch_func, **kwargs):
   return functools.partial(register_torch_function_op, torch_func, **kwargs)
@@ -280,7 +290,7 @@ def empty(*size: Sequence[int], dtype=None, **kwargs):
 
 @register_function(torch.arange, is_jax_function=False)
 def arange(
-  start, end=None, step=None, 
+  start, end=None, step=None,
   out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False,
   pin_memory=None,
 ):
@@ -414,3 +424,43 @@ def linalg_tensorsolve(A, b, dims=None):
   if A.shape[:b.ndim] != b.shape:
     b = jnp.reshape(b, A.shape[:b.ndim])
   return jnp.linalg.tensorsolve(A, b, axes=dims)
+
+@register_function(torch.linspace)
+def linspace(start, end, steps=100, *, dtype=None, device=None, **kwargs):
+    env = kwargs.get("env")
+
+    if dtype is None:
+        dtype = torch.float32
+
+    jdtype = TORCH_TO_JAX_DTYPE.get(dtype)
+    if jdtype is None:
+        raise RuntimeError(f"Attempting to convert unknown type: {dtype} to jax type")
+
+    start_j = jnp.array(start, dtype=jdtype)
+    end_j   = jnp.array(end,   dtype=jdtype)
+
+    result = jnp.linspace(start_j, end_j, num=steps, dtype=jdtype)
+    return torch_xla2.tensor.XLATensor2(result, env)
+
+
+@register_function(torch.logspace)
+def logspace(start, end, steps=100, base=10.0, *, dtype=None, device=None, **kwargs):
+    """
+    Generates a sequence of numbers spaced evenly on a log scale, from
+    base**start to base**end in `steps` points.
+    """
+    env = kwargs.get("env")
+
+    if dtype is None:
+        dtype = torch.float32
+
+    jdtype = TORCH_TO_JAX_DTYPE.get(dtype)
+    if jdtype is None:
+        raise RuntimeError(f"Attempting to convert unknown type: {dtype} to jax type")
+
+    start_j = jnp.array(start, dtype=jdtype)
+    end_j   = jnp.array(end,   dtype=jdtype)
+
+    result = jnp.logspace(start_j, end_j, num=steps, base=base, dtype=jdtype)
+
+    return torch_xla2.tensor.XLATensor2(result, env)