Try to work arround AssertionError in XLA

Signed-off-by: Fabrice Normandin <normandf@mila.quebec>

Author: lebrice

torch_jax_interop/to_jax.py

@@ -93,8 +93,14 @@ def _direct_conversion(v: torch.Tensor) -> jax.Array:
     return jax_from_dlpack(v, copy=False)
 
 
-def _to_from_dlpack(v: torch.Tensor, ignore_deprecation_warning: bool = True) -> jax.Array:
-    with warnings.catch_warnings() if ignore_deprecation_warning else contextlib.nullcontext():
+def _to_from_dlpack(
+    v: torch.Tensor, ignore_deprecation_warning: bool = True
+) -> jax.Array:
+    with (
+        warnings.catch_warnings()
+        if ignore_deprecation_warning
+        else contextlib.nullcontext()
+    ):
         # Only way to get this to work for CPU seems to be with to/from dlpack... so we have to use this deprecated
         # conversion method for now.
         # todo: Should we let it though though?
@@ -136,7 +142,29 @@ def torch_to_jax_tensor(value: torch.Tensor) -> jax.Array:
             return _direct_conversion(value.flatten()).reshape(value.shape)
 
     try:
-        return _direct_conversion(value)
+        # Try using the "new" way to convert using from_dlpack directly
+        return jax_from_dlpack(
+            value, device=torch_to_jax_device(value.device), copy=None
+        )
+    except AssertionError as err:
+        if not err.args[0].startswith("Unexpected XLA layout override"):
+            raise
+        # Some "AssertionError: Unexpected XLA layout override"
+        # Try using the "old" way to convert using from_dlpack of a dlpack tensor.
+        try:
+            dlpack = torch_to_dlpack(value)
+            return jax_from_dlpack(dlpack, copy=False)
+        except jaxlib.xla_extension.XlaRuntimeError as err:
+            log_once(
+                logger,
+                message=(
+                    f"Unable to view tensor of shape {tuple(value.shape)} as a jax.Array in-place:\n"
+                    f"'{err}'\n"
+                    f"Tensors of this shape will be flattened and unflattened (which may or "
+                    f"may not involve making a copy of the tensor's data)."
+                ),
+                level=logging.WARNING,
+            )
     except jaxlib.xla_extension.XlaRuntimeError as err:
         log_once(
             logger,

torch_jax_interop/utils.py

@@ -11,6 +11,7 @@ def log_once(logger: logging.Logger, message: str, level: int):
     logger.log(level=level, msg=message, stacklevel=2)
 
 
+# NOTE: Done like this to preserve the original function signature.
 log_once = functools.cache(log_once)