Repository: google/jax
Issue: bitcast_convert_type does not support cross-size casts

### Description

The [current definition](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.bitcast_convert_type.html) of `jax.lax.bitcast_convert_type` requires that the source and destination types have the same size. The [underlying XLA op](https://www.tensorflow.org/xla/operation_semantics#bitcastconverttype) does not have this requirement and has well-defined semantics for differently-sized types.

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

_No response_

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

TPU

### Additional system info

_No response_

### NVIDIA GPU info

_No response_

Changed files:
jax/_src/lax/lax.py, tests/lax_test.py

diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
index 8d5d4ecc1..5347aba6f 100644
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -595,16 +595,26 @@ def bitcast_convert_type(operand: ArrayLike, new_dtype: DTypeLike) -> Array:
 
   Wraps XLA's `BitcastConvertType
   <https://www.tensorflow.org/xla/operation_semantics#bitcastconverttype>`_
-  operator, which performs a bit cast from one type to another. The bitwidth
-  of the source and destination types must match.
+  operator, which performs a bit cast from one type to another.
+
+  The output shape depends on the size of the input and output dtypes with
+  the following logic::
+
+    if new_dtype.itemsize == operand.dtype.itemsize:
+      output_shape = operand.shape
+    if new_dtype.itemsize < operand.dtype.itemsize:
+      output_shape = (*operand.shape, operand.dtype.itemsize // new_dtype.itemsize)
+    if new_dtype.itemsize > operand.dtype.itemsize:
+      assert operand.shape[-1] * operand.dtype.itemsize == new_dtype.itemsize
+      output_shape = operand.shape[:-1]
 
   Args:
     operand: an array or scalar value to be cast
     new_dtype: the new type. Should be a NumPy type.
 
   Returns:
-    An array with the same shape as `operand`, bitcast elementwise to
-    `new_dtype`.
+    An array of shape `output_shape` (see above) and type `new_dtype`,
+    constructed from the same bits as operand.
   """
   new_dtype = dtypes.canonicalize_dtype(new_dtype)
   return bitcast_convert_type_p.bind(operand, new_dtype=new_dtype)
@@ -2418,15 +2428,26 @@ mlir.register_lowering(convert_element_type_p, _convert_element_type_lower)
 
 
 def _bitcast_convert_type_shape_rule(operand, *, new_dtype):
-  return operand.shape
+  old_dtype = dtypes.canonicalize_dtype(operand.dtype)
+  new_dtype = dtypes.canonicalize_dtype(new_dtype)
+
+  if old_dtype.itemsize == new_dtype.itemsize:
+    return operand.shape
+  elif old_dtype.itemsize > new_dtype.itemsize:
+    return (*operand.shape, old_dtype.itemsize // new_dtype.itemsize)
+  elif operand.shape[-1] * old_dtype.itemsize != new_dtype.itemsize:
+    raise ValueError(
+      f"Attempting to convert array of shape {operand.shape} "
+      f"from {str(old_dtype)} of size {old_dtype.itemsize} "
+      f"to {str(new_dtype)} of size {new_dtype.itemsize}, "
+      f"but {operand.shape[-1]} * {old_dtype.itemsize} != {new_dtype.itemsize}")
+  return operand.shape[:-1]
 
 def _bitcast_convert_type_dtype_rule(operand, *, new_dtype):
   old_dtype = dtypes.canonicalize_dtype(operand.dtype)
   if dtypes.issubdtype(old_dtype, np.bool_) or dtypes.issubdtype(old_dtype, np.complexfloating):
     if old_dtype != new_dtype:
       raise TypeError(f"`bitcast_convert_type` for operand type ({old_dtype}) cannot have different destination type ({new_dtype})")
-  if np.dtype(old_dtype).itemsize != np.dtype(new_dtype).itemsize:
-    raise TypeError(f"`bitcast_convert_type` for operand type ({old_dtype}) must have destination type ({new_dtype}) of same size.")
   return new_dtype
 
 bitcast_convert_type_p = standard_primitive(
diff --git a/tests/lax_test.py b/tests/lax_test.py
index 62acb2b0b..57a3dfda6 100644
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -281,16 +281,32 @@ class LaxTest(jtu.JaxTestCase):
     self._CheckAgainstNumpy(numpy_op, op, args_maker)
 
   @jtu.sample_product(
-    [dict(from_dtype=from_dtype, to_dtype=to_dtype)
-     for from_dtype, to_dtype in itertools.product(
-      [np.float32, np.int32, "float32", "int32"], repeat=2)],
+    from_dtype=jtu.dtypes.all_floating + jtu.dtypes.all_integer + jtu.dtypes.all_unsigned,
+    to_dtype=jtu.dtypes.all_floating + jtu.dtypes.all_integer + jtu.dtypes.all_unsigned,
+    shape = [(), (2,), (2, 3)]
   )
-  def testBitcastConvertType(self, from_dtype, to_dtype):
+  def testBitcastConvertType(self, from_dtype, to_dtype, shape):
     rng = jtu.rand_default(self.rng())
-    args_maker = lambda: [rng((2, 3), from_dtype)]
+    itemsize_in = np.dtype(from_dtype).itemsize
+    itemsize_out = np.dtype(to_dtype).itemsize
+    if itemsize_in < itemsize_out:
+      shape = (*shape, itemsize_out // itemsize_in)
+    args_maker = lambda: [rng(shape, from_dtype)]
     op = lambda x: lax.bitcast_convert_type(x, to_dtype)
     self._CompileAndCheck(op, args_maker)
 
+    # Test the shape and dtype of the output. We avoid testing the values here
+    # because the bitwise representation may vary from platform to platform.
+    out = op(*args_maker())
+    if itemsize_in == itemsize_out:
+      expected_shape = shape
+    elif itemsize_in < itemsize_out:
+      expected_shape = shape[:-1]
+    else:
+      expected_shape = (*shape, itemsize_in // itemsize_out)
+    self.assertEqual(out.dtype, to_dtype)
+    self.assertEqual(out.shape, expected_shape)
+
   @jtu.sample_product(
     [dict(from_dtype=from_dtype, to_dtype=to_dtype)
      for from_dtype, to_dtype in itertools.product(