# Get available devices
print("Available Devices:", jax.devices())

# Check if GPU is available
if jax.default_backend() == "gpu":
    print("Using GPU:", jax.devices("gpu"))
elif jax.default_backend() == "tpu":
    print("Using TPU:", jax.devices("tpu"))
else:
    print("Using CPU")

        __
        __