Repository: google/jax

Issue:
BUG: can create 64-bit arrays using dlpack when x64 is disabled
### Description

Repro:
```python
import tensorflow as tf
import jax
import jax.dlpack

assert not jax.config.jax_enable_x64
x = jax.dlpack.from_dlpack(tf.experimental.dlpack.to_dlpack(tf.ones(5, tf.int64)))
assert x.dtype == jax.numpy.int64
```

This can cause loud and confusing errors down the line when our MHLO expects 32-bit numbers but gets 64-bit. For example, with our illegal 64-bit `x` array, try doing `x[0]`.

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

HEAD/0.3.15

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

CPU

### Additional System Info

_No response_


Discussion:
> This can cause loud and confusing errors down the line

Arguably a separate issue to file regarding this: if we do somehow end up with unexpected 64-bit device (or abstract) arrays, we ideally catch that and err early on, ahead of lowering in particular.
Reviving this issue.

Hi, I had experienced the same exact problem and the error I got was:
```python
XlaRuntimeError: UNKNOWN: <unknown>:0: error: type of return operand 0 ('tensor<1xi64>') doesn't match function result type ('tensor<1xi32>') in function @main
<unknown>:0: note: see current operation: "func.return"(%0) : (tensor<1xi64>) -> ()
<unknown>:0: note: in bytecode version 0 produced by: MLIR16.0.0git
```
Which wasn't very intuitive to figure out. Furthermore, it seems that some tensorflow_datasets may use int64 to describe labels (like in tfds.load('mnist')) and if used together with dlpack it will result in this issue.

Isn´t there a way to protect users from this, like a more intuitive RuntimeError, Warning?

For context I will leave [here the colab](https://colab.research.google.com/drive/1dXdlgdduOQ9SY_MIahMtpOZk1E5WcmLb?usp=sharing) that I used to replicate this issue.
@hawkinsp @jakevdp – So long as the x64 flag exists and means what it means, how should we behave on `jax.dlpack.from_dlpack` given a 64-bit-dtyped array, when x64 mode is off? Some options:

1. Error
2. Convert to 32-bit silently
3. Convert to 32-bit and `warnings.warn`

I'll send a PR for what we think is best.
converting silently is consistent with how the x64 flag works in similar situations, e.g. `jnp.asarray(numpy_arr)`

Resolving patch:
diff --git a/jax/_src/dlpack.py b/jax/_src/dlpack.py
--- a/jax/_src/dlpack.py
+++ b/jax/_src/dlpack.py
@@ -70,4 +70,5 @@ def from_dlpack(dlpack):
   xla_shape = buf.xla_shape()
   assert not xla_shape.is_tuple()
   aval = core.ShapedArray(xla_shape.dimensions(), xla_shape.numpy_dtype())
-  return dispatch.maybe_create_array_from_da(buf, aval, buf.device())
+  return jnp.asarray(           # asarray ensures dtype canonicalization
+      dispatch.maybe_create_array_from_da(buf, aval, buf.device()))


Test patch:
diff --git a/tests/array_interoperability_test.py b/tests/array_interoperability_test.py
--- a/tests/array_interoperability_test.py
+++ b/tests/array_interoperability_test.py
@@ -137,6 +137,14 @@ def testJaxToTensorFlow(self, shape, dtype):
     y = tf.experimental.dlpack.from_dlpack(dlpack)
     self.assertAllClose(np, y.numpy())
 
+  @unittest.skipIf(not tf, "Test requires TensorFlow")
+  def testTensorFlowToJaxInt64(self):
+    # See https://github.com/google/jax/issues/11895
+    x = jax.dlpack.from_dlpack(
+        tf.experimental.dlpack.to_dlpack(tf.ones((2, 3), tf.int64)))
+    dtype_expected = jnp.int64 if config.x64_enabled else jnp.int32
+    self.assertEqual(x.dtype, dtype_expected)
+
   @jtu.sample_product(
     shape=all_shapes,
     dtype=torch_dtypes,
@@ -182,6 +190,14 @@ def testJaxToTorch(self, shape, dtype):
     y = torch.utils.dlpack.from_dlpack(dlpack)
     self.assertAllClose(np, y.cpu().numpy())
 
+  @unittest.skipIf(not torch, "Test requires PyTorch")
+  def testTorchToJaxInt64(self):
+    # See https://github.com/google/jax/issues/11895
+    x = jax.dlpack.from_dlpack(
+        torch.utils.dlpack.to_dlpack(torch.ones((2, 3), dtype=torch.int64)))
+    dtype_expected = jnp.int64 if config.x64_enabled else jnp.int32
+    self.assertEqual(x.dtype, dtype_expected)
+
   @jtu.sample_product(
     shape=all_shapes,
     dtype=torch_dtypes,