Repository: google/jax

Issue:
`broadcast_to` raise Disallowed host-to-device transfer when used inside jax.jit
### Description

Reproduction:

```python
@jax.jit
def test():
  return jnp.broadcast_to(0.1, (4,))

with jax.transfer_guard('disallow'):
  test()
```
Fail with:
```
XlaRuntimeError: INVALID_ARGUMENT: Disallowed host-to-device transfer: aval=ShapedArray(float32[]), dst_sharding=SingleDeviceSharding(device=TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0))
```

The `0.1` is a constant in the XLA graph, so this should not trigger any host-to-device transfer.

It's possible to silence the issue with `jnp.broadcast_to(jnp.asarray(0.1), (4,))` but original issue is a bug.

### What jax/jaxlib version are you using?

HEAD

### Which accelerator(s) are you using?

TPU



Discussion:
Thanks for the report. It looks like the issue comes from the implementation of `lax.asarray` (used by `broadcast_to`), which uses `device_put` for non-jax-array inputs: https://github.com/google/jax/blob/31efc2cc6ae6f7f20b78c742def23afcd6507625/jax/_src/lax/lax.py#L129-L137

I think we need to use a different strategy for converting scalars to arrays if we want to avoid triggering the transfer guard here.

Here's the more direct repro:
```python
import jax

with jax.transfer_guard('disallow'):
  jax.jit(lambda: jax.numpy.asarray(0.1))()  # no error
  jax.jit(lambda: jax._src.lax.lax.asarray(0.1))()  # error
```

Resolving patch:
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -132,9 +132,8 @@ def asarray(x: ArrayLike) -> Array:
   """Lightweight conversion of ArrayLike input to Array output."""
   if isinstance(x, Array):
     return x
-  if isinstance(x, np.ndarray) or np.isscalar(x):
-    # Call device_put_impl directly to avoid binding the primitive.
-    return dispatch._device_put_impl(x)
+  if isinstance(x, (np.ndarray, np.generic, bool, int, float, builtins.complex)):
+    return _convert_element_type(x, weak_type=dtypes.is_weakly_typed(x))
   else:
     raise TypeError(f"asarray: expected ArrayLike, got {x} of type {type(x)}.")
 


Test patch:
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -2726,6 +2726,28 @@ def _step(carry, arg):
 
     a, b = jax.lax.scan(_step, 0, jnp.arange(4, dtype=jnp.complex64))
 
+  @parameterized.parameters([float, np.array, np.float32, jnp.float32])
+  def testAsarray(self, typ):
+    x = typ(1.0)
+    x_arr = lax_internal.asarray(x)
+    self.assertArraysEqual(x, x_arr)
+    self.assertIsInstance(x_arr, jax.Array)
+
+    # jaxpr should not bind any primitives, whether called directly or
+    # as a closure:
+    jaxpr = jax.make_jaxpr(lax_internal.asarray)(x)
+    self.assertLen(jaxpr.eqns, 0)
+
+    asarray_closure = lambda: lax_internal.asarray(x)
+    jaxpr = jax.make_jaxpr(asarray_closure)()
+    self.assertLen(jaxpr.eqns, 0)
+
+    # Regression test for https://github.com/google/jax/issues/19334
+    # lax.asarray as a closure should not trigger transfer guard.
+    with jax.transfer_guard('disallow'):
+      jax.jit(asarray_closure)()
+
+
 
 class LazyConstantTest(jtu.JaxTestCase):
   def _Check(self, make_const, expected):