device_put leaks references on numpy inputs
First reported in #14876; here's a minimal repro that reproduces on CPU and GPU:
```python
import sys
import jax
import numpy as np

x = np.arange(1000, dtype='float32')
print(sys.getrefcount(x))
# 2

y = jax.device_put(x)
print(sys.getrefcount(x))
# 3

del y
print(sys.getrefcount(x))
# 3
```
diff --git a/jax/_src/lib/__init__.py b/jax/_src/lib/__init__.py
index 6c434d53c..3144402c8 100644
--- a/jax/_src/lib/__init__.py
+++ b/jax/_src/lib/__init__.py
@@ -15,10 +15,9 @@
 # This module is largely a wrapper around `jaxlib` that performs version
 # checking on import.
 
-import platform
+import gc
 import re
 import os
-import warnings
 from typing import Optional, Tuple
 
 try:
@@ -94,6 +93,11 @@ pytree = xla_client._xla.pytree
 jax_jit = xla_client._xla.jax_jit
 pmap_lib = xla_client._xla.pmap_lib
 
+# XLA garbage collection: see https://github.com/google/jax/issues/14882
+def _xla_gc_callback(*args):
+  xla_client._xla.collect_garbage()
+gc.callbacks.append(_xla_gc_callback)
+
 import jaxlib.gpu_solver as gpu_solver  # pytype: disable=import-error
 import jaxlib.gpu_sparse as gpu_sparse  # pytype: disable=import-error
 import jaxlib.gpu_prng as gpu_prng  # pytype: disable=import-error
diff --git a/tests/api_test.py b/tests/api_test.py
index 978c22547..18e525500 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -9629,5 +9629,18 @@ class AutodidaxTest(jtu.JaxTestCase):
     autodidax_module = importlib.util.module_from_spec(spec)
     spec.loader.exec_module(autodidax_module)
 
+class GarbageCollectionTest(jtu.JaxTestCase):
+  def test_xla_gc_callback(self):
+    # https://github.com/google/jax/issues/14882
+    x_np = np.arange(10, dtype='int32')
+    x_jax = jax.device_put(x_np)
+    x_np_weakref = weakref.ref(x_np)
+
+    del x_np
+    del x_jax
+    gc.collect()
+
+    assert x_np_weakref() is None
+
 if __name__ == '__main__':
   absltest.main(testLoader=jtu.JaxTestLoader())
['jax/_src/lib/__init__.py', 'tests/api_test.py']