Change of Variable in HMC

Change of Variable in HMC#

Rat tumor problem: We have J certain kinds of rat tumor diseases. For each kind of tumor, we test \(N_{j}\) people/animals and among those \(y_{j}\) tested positive. Here we assume that \(y_{j}\) is distrubuted with Binom(\(N_{j}\), \(\theta_{j}\)). Our objective is to approximate \(\theta_{j}\) for each type of tumor.

In particular we use following binomial hierarchical model where \(y_{j}\) and \(N_{j}\) are observed variables.

\[\begin{split}\begin{align} y_{j} &\sim \text{Binom}(N_{j}, \theta_{j}) \label{eq:1} \\ \theta_{j} &\sim \text{Beta}(a, b) \label{eq:2} \\ p(a, b) &\propto (a+b)^{-5/2} \end{align}\end{split}\]

Hide code cell content

 
import matplotlib.pyplot as plt

plt.rcParams["axes.spines.right"] = False
plt.rcParams["axes.spines.top"] = False
plt.rcParams["xtick.labelsize"] = 12
plt.rcParams["ytick.labelsize"] = 12

import pandas as pd

pd.set_option("display.max_rows", 80)

import jax

from datetime import date
rng_key = jax.random.key(int(date.today().strftime("%Y%m%d")))

Hide code cell content

 
import arviz as az
import jax.numpy as jnp

import blackjax
import tensorflow_probability.substrates.jax as tfp

tfd = tfp.distributions
tfb = tfp.bijectors

az.rcParams["plot.max_subplots"] = 200

Hide code cell content

 
# index of array is type of tumor and value shows number of total people tested.
group_size = jnp.array(
    [
        20,
        20,
        20,
        20,
        20,
        20,
        20,
        19,
        19,
        19,
        19,
        18,
        18,
        17,
        20,
        20,
        20,
        20,
        19,
        19,
        18,
        18,
        25,
        24,
        23,
        20,
        20,
        20,
        20,
        20,
        20,
        10,
        49,
        19,
        46,
        27,
        17,
        49,
        47,
        20,
        20,
        13,
        48,
        50,
        20,
        20,
        20,
        20,
        20,
        20,
        20,
        48,
        19,
        19,
        19,
        22,
        46,
        49,
        20,
        20,
        23,
        19,
        22,
        20,
        20,
        20,
        52,
        46,
        47,
        24,
        14,
    ],
    dtype=jnp.float32,
)

# index of array is type of tumor and value shows number of positve people.
n_of_positives = jnp.array(
    [
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        1,
        2,
        2,
        2,
        2,
        2,
        2,
        2,
        2,
        2,
        1,
        5,
        2,
        5,
        3,
        2,
        7,
        7,
        3,
        3,
        2,
        9,
        10,
        4,
        4,
        4,
        4,
        4,
        4,
        4,
        10,
        4,
        4,
        4,
        5,
        11,
        12,
        5,
        5,
        6,
        5,
        6,
        6,
        6,
        6,
        16,
        15,
        15,
        9,
        4,
    ],
    dtype=jnp.float32,
)

# number of different kind of rat tumors
n_rat_tumors = len(group_size)

Hide code cell source

 
_, axes = plt.subplots(2, 1, figsize=(12, 6))
axes[0].bar(range(n_rat_tumors), n_of_positives)
axes[0].set_title("No. of positives for each tumor type", fontsize=14);
axes[1].bar(range(n_rat_tumors), group_size)
axes[1].set_xlabel("tumor type", fontsize=12)
axes[1].set_title("Group size for each tumor type", fontsize=14)
plt.tight_layout();
../_images/0e167109cc355367c9cc15263a18f002a2bd89e639eea4281f78cd65af51e064.png

Posterior Sampling#

Now we use Blackjax’s NUTS algorithm to get posterior samples of \(a\), \(b\), and \(\theta\)

from collections import namedtuple

params = namedtuple("model_params", ["a", "b", "thetas"])


def joint_logdensity(params):
    # improper prior for a,b
    logdensity_ab = jnp.log(jnp.power(params.a + params.b, -2.5))

    # logdensity prior of theta
    logdensity_thetas = tfd.Beta(params.a, params.b).log_prob(params.thetas).sum()

    # loglikelihood of y
    logdensity_y = jnp.sum(
        tfd.Binomial(group_size, probs=params.thetas).log_prob(n_of_positives)
    )

    return logdensity_ab + logdensity_thetas + logdensity_y

We take initial parameters from uniform distribution

rng_key, init_key = jax.random.split(rng_key)
n_params = n_rat_tumors + 2


def init_param_fn(seed):
    """
    initialize a, b & thetas
    """
    key1, key2, key3 = jax.random.split(seed, 3)
    return params(
        a=tfd.Uniform(0, 3).sample(seed=key1),
        b=tfd.Uniform(0, 3).sample(seed=key2),
        thetas=tfd.Uniform(0, 1).sample(n_rat_tumors, seed=key3),
    )


init_param = init_param_fn(init_key)
joint_logdensity(init_param)  # sanity check

Array(-1270.7018, dtype=float32)

Now we use blackjax’s window adaption algorithm to get NUTS kernel and initial states. Window adaption algorithm will automatically configure inverse_mass_matrix and step size

%%time
warmup = blackjax.window_adaptation(blackjax.nuts, joint_logdensity)

# we use 4 chains for sampling
n_chains = 4
rng_key, init_key, warmup_key = jax.random.split(rng_key, 3)
init_keys = jax.random.split(init_key, n_chains)
init_params = jax.vmap(init_param_fn)(init_keys)

@jax.vmap
def call_warmup(seed, param):
    (initial_states, tuned_params), _ = warmup.run(seed, param, 1000)
    return initial_states, tuned_params

warmup_keys = jax.random.split(warmup_key, n_chains)
initial_states, tuned_params = jax.jit(call_warmup)(warmup_keys, init_params)

CPU times: user 9.5 s, sys: 784 ms, total: 10.3 s
Wall time: 5.64 s

Now we write inference loop for multiple chains

def inference_loop_multiple_chains(
    rng_key, initial_states, tuned_params, log_prob_fn, num_samples, num_chains
):
    kernel = blackjax.nuts.build_kernel()

    def step_fn(key, state, **params):
        return kernel(key, state, log_prob_fn, **params)

    def one_step(states, rng_key):
        keys = jax.random.split(rng_key, num_chains)
        states, infos = jax.vmap(step_fn)(keys, states, **tuned_params)
        return states, (states, infos)

    keys = jax.random.split(rng_key, num_samples)
    _, (states, infos) = jax.lax.scan(one_step, initial_states, keys)

    return (states, infos)

%%time
n_samples = 1000
rng_key, sample_key = jax.random.split(rng_key)
states, infos = inference_loop_multiple_chains(
    sample_key, initial_states, tuned_params, joint_logdensity, n_samples, n_chains
)

CPU times: user 8.83 s, sys: 389 ms, total: 9.22 s
Wall time: 3.92 s

Arviz Plots#

We have all our posterior samples stored in states.position dictionary and infos store additional information like acceptance probability, divergence, etc. Now, we can use certain diagnostics to judge if our MCMC samples are converged on stationary distribution. Some of widely diagnostics are trace plots, potential scale reduction factor (R hat), divergences, etc. Arviz library provides quicker ways to anaylze these diagnostics. We can use arviz.summary() and arviz_plot_trace(), but these functions take specific format (arviz’s trace) as a input.

Hide code cell content

 
def arviz_trace_from_states(states, info, burn_in=0):
    position = states.position
    if isinstance(position, jax.Array):  # if states.position is array of samples
        position = dict(samples=position)
    else:
        try:
            position = position._asdict()
        except AttributeError:
            pass

    samples = {}
    for param in position.keys():
        ndims = len(position[param].shape)
        if ndims >= 2:
            samples[param] = jnp.swapaxes(position[param], 0, 1)[
                :, burn_in:
            ]  # swap n_samples and n_chains
            divergence = jnp.swapaxes(info.is_divergent[burn_in:], 0, 1)

        if ndims == 1:
            divergence = info.is_divergent
            samples[param] = position[param]

    trace = az.from_dict({"posterior": samples, "sample_stats": {"diverging": divergence}})
    return trace

# make arviz trace from states
trace = arviz_trace_from_states(states, infos)
summ_df = az.summary(trace)
summ_df
mean sd eti89_lb eti89_ub ess_bulk ess_tail r_hat mcse_mean mcse_sd
a 0.69 0.07 0.61 0.83 4 6 2.43 0.035 0.019
b 2.3 0.7 1.3 3.8 4 5 3.38 0.37 0.21
thetas[0] 0.05 0.06 0.0006 0.15 4 7 2.53 0.027 0.015
thetas[1] 0.03 0.015 0.0051 0.051 6 8 1.67 0.0059 0.0041
thetas[2] 0.1 0.1 0.0027 0.27 6 9 1.91 0.051 0.029
thetas[3] 0.023 0.02 0.00056 0.054 18 14 1.15 0.0038 0.0052
thetas[4] 0.1 0.17 0.0012 0.41 5 13 2.23 0.083 0.047
thetas[5] 0.05 0.04 0.0014 0.12 5 8 1.93 0.016 0.01
thetas[6] 0.08 0.04 0.0077 0.13 4 6 2.42 0.018 0.0092
thetas[7] 0.04 0.04 0.002 0.13 5 4 2.01 0.02 0.014
thetas[8] 0.1 0.12 0.002 0.33 5 4 1.93 0.058 0.039
thetas[9] 0.04 0.05 0.00024 0.14 6 10 1.77 0.024 0.014
thetas[10] 0.02 0.02 0.0003 0.054 9 24 1.39 0.0059 0.0058
thetas[11] 0.02 0.03 0.0002 0.08 6 7 1.68 0.012 0.012
thetas[12] 0.1 0.11 0.0023 0.31 4 7 3.34 0.053 0.027
thetas[13] 0.2 0.15 0.0064 0.46 4 5 2.44 0.076 0.035
thetas[14] 0.09 0.05 0.028 0.18 5 9 1.97 0.023 0.013
thetas[15] 0.07 0.05 0.006 0.16 5 4 2.27 0.023 0.014
thetas[16] 0.1 0.06 0.022 0.19 6 7 1.98 0.025 0.0097
thetas[17] 0.2 0.2 0.012 0.54 4 6 3.21 0.098 0.05
thetas[18] 0.13 0.09 0.034 0.27 5 4 2.02 0.041 0.021
thetas[19] 0.09 0.09 0.032 0.33 6 4 1.65 0.043 0.047
thetas[20] 0.1 0.1 0.022 0.47 4 4 3.47 0.068 0.054
thetas[21] 0.2 0.14 0.042 0.42 4 5 2.54 0.07 0.035
thetas[22] 0.09 0.04 0.042 0.15 7 14 1.48 0.015 0.01
thetas[23] 0.2 0.12 0.069 0.41 4 7 2.45 0.058 0.037
thetas[24] 0.08 0.03 0.041 0.12 9 17 1.38 0.0086 0.0065
thetas[25] 0.15 0.07 0.062 0.25 4 9 2.52 0.032 0.012
thetas[26] 0.16 0.09 0.016 0.29 4 4 3.14 0.046 0.022
thetas[27] 0.2 0.17 0.053 0.52 4 8 3.19 0.085 0.045
thetas[28] 0.3 0.2 0.053 0.63 4 4 3.49 0.1 0.044
thetas[29] 0.14 0.09 0.029 0.27 4 5 2.54 0.042 0.018
thetas[30] 0.11 0.06 0.025 0.2 4 6 3.05 0.032 0.0096
thetas[31] 0.2 0.13 0.057 0.39 5 13 2.38 0.066 0.017
thetas[32] 0.1 0.03 0.061 0.15 6 20 1.66 0.012 0.0072
thetas[33] 0.3 0.2 0.043 0.64 4 4 3.37 0.11 0.05
thetas[34] 0.13 0.07 0.057 0.26 4 4 2.65 0.03 0.019
thetas[35] 0.12 0.04 0.059 0.18 4 9 2.37 0.022 0.0092
thetas[36] 0.3 0.16 0.12 0.53 4 6 3.46 0.079 0.037
thetas[37] 0.16 0.08 0.078 0.29 5 10 2.35 0.036 0.018
thetas[38] 0.13 0.03 0.1 0.19 5 18 1.80 0.013 0.012
thetas[39] 0.14 0.05 0.089 0.23 4 4 2.76 0.023 0.0097
thetas[40] 0.2 0.11 0.099 0.39 4 4 3.16 0.053 0.031
thetas[41] 0.16 0.08 0.069 0.28 5 10 2.15 0.038 0.015
thetas[42] 0.19 0.03 0.15 0.26 5 11 2.00 0.014 0.011
thetas[43] 0.21 0.05 0.16 0.31 5 6 1.99 0.026 0.014
thetas[44] 0.2 0.14 0.073 0.47 4 8 3.16 0.069 0.032
thetas[45] 0.3 0.13 0.11 0.46 4 9 2.43 0.064 0.029
thetas[46] 0.23 0.04 0.15 0.3 4 4 2.89 0.021 0.012
thetas[47] 0.2 0.12 0.09 0.4 4 6 2.63 0.056 0.02
thetas[48] 0.3 0.15 0.12 0.58 5 5 2.26 0.072 0.044
thetas[49] 0.2 0.1 0.059 0.35 4 11 2.75 0.051 0.018
thetas[50] 0.22 0.07 0.12 0.36 5 4 2.31 0.032 0.024
thetas[51] 0.23 0.05 0.15 0.31 5 7 1.93 0.022 0.014
thetas[52] 0.19 0.07 0.091 0.33 4 5 3.10 0.033 0.029
thetas[53] 0.2 0.08 0.11 0.32 5 4 1.87 0.034 0.015
thetas[54] 0.39 0.04 0.31 0.44 5 6 2.07 0.019 0.014
thetas[55] 0.23 0.04 0.12 0.28 12 16 1.54 0.0099 0.011
thetas[56] 0.28 0.08 0.18 0.41 5 6 2.28 0.037 0.017
thetas[57] 0.28 0.05 0.2 0.35 5 7 1.84 0.02 0.015
thetas[58] 0.42 0.03 0.36 0.48 5 4 1.86 0.014 0.01
thetas[59] 0.3 0.12 0.16 0.5 4 4 2.97 0.057 0.032
thetas[60] 0.3 0.2 0.17 0.66 4 7 2.83 0.095 0.052
thetas[61] 0.34 0.09 0.18 0.46 4 8 2.91 0.041 0.029
thetas[62] 0.29 0.07 0.2 0.43 5 5 1.88 0.03 0.02
thetas[63] 0.24 0.04 0.17 0.3 5 7 1.84 0.019 0.0098
thetas[64] 0.3 0.15 0.14 0.54 4 7 2.97 0.072 0.033
thetas[65] 0.24 0.08 0.15 0.37 4 6 2.41 0.04 0.014
thetas[66] 0.32 0.07 0.21 0.42 4 5 3.61 0.034 0.012
thetas[67] 0.28 0.04 0.2 0.35 4 9 2.37 0.02 0.012
thetas[68] 0.32 0.08 0.2 0.44 4 4 2.91 0.041 0.018
thetas[69] 0.48 0.08 0.37 0.58 4 5 2.41 0.037 0.013
thetas[70] 0.4 0.22 0.094 0.71 4 8 2.85 0.11 0.047

r_hat is showing measure of each chain is converged to stationary distribution. r_hat should be less than or equal to 1.01, here we get r_hat far from 1.01 for each latent sample.

Hide code cell source

 
az.plot_trace(trace)
plt.tight_layout();

/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)

/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)

/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)

/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)

/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)

/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)

/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)

/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)

/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)

/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)

/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)

/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)

/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)

/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)

/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)

/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)

/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)

/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)

/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)

/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)

/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)

/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)

/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)

/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)

/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)

/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)

/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
/opt/hostedtoolcache/Python/3.13.12/x64/lib/python3.13/site-packages/xarray/core/duck_array_ops.py:265: UserWarning: Explicitly requested dtype int64 requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  return xp.astype(data, dtype, **kwargs)
../_images/92b3b3e7845793e65f62ecbb7fad29e04ef184211fc71a44f5732920dd32004e.png

Trace plots also looks terrible and does not seems to be converged! Also, black band shows that every sample is diverged from original distribution. So what’s wrong happeing here?

Well, it’s related to support of latent variable. In HMC, the latent variable must be in an unconstrained space, but in above model theta is constrained in between 0 to 1. We can use change of variable trick to solve above problem

Change of Variable#

We can sample from logits which is in unconstrained space and in joint_logdensity() we can convert logits to theta by suitable bijector (sigmoid). We calculate jacobian (first order derivaive) of bijector to tranform one probability distribution to another

transform_fn = jax.nn.sigmoid
log_jacobian_fn = lambda logit: jnp.log(jnp.abs(jnp.diag(jax.jacfwd(transform_fn)(logit))))

Alternatively, using the bijector class in TFP directly:

bij = tfb.Sigmoid()
transform_fn = bij.forward
log_jacobian_fn = bij.forward_log_det_jacobian

params = namedtuple("model_params", ["a", "b", "logits"])

def joint_logdensity_change_of_var(params):
    # change of variable
    thetas = transform_fn(params.logits)
    log_det_jacob = jnp.sum(log_jacobian_fn(params.logits))

    # improper prior for a,b
    logdensity_ab = jnp.log(jnp.power(params.a + params.b, -2.5))

    # logdensity prior of theta
    logdensity_thetas = tfd.Beta(params.a, params.b).log_prob(thetas).sum()

    # loglikelihood of y
    logdensity_y = jnp.sum(
        tfd.Binomial(group_size, probs=thetas).log_prob(n_of_positives)
    )

    return logdensity_ab + logdensity_thetas + logdensity_y + log_det_jacob

except for the change of variable in joint_logdensity() function, everthing will remain same

rng_key, init_key = jax.random.split(rng_key)


def init_param_fn(seed):
    """
    initialize a, b & logits
    """
    key1, key2, key3 = jax.random.split(seed, 3)
    return params(
        a=tfd.Uniform(0, 3).sample(seed=key1),
        b=tfd.Uniform(0, 3).sample(seed=key2),
        logits=tfd.Uniform(-2, 2).sample(n_rat_tumors, key3),
    )


init_param = init_param_fn(init_key)
joint_logdensity_change_of_var(init_param)  # sanity check

Array(-907.1453, dtype=float32)

%%time
warmup = blackjax.window_adaptation(blackjax.nuts, joint_logdensity_change_of_var)

# we use 4 chains for sampling
n_chains = 4
rng_key, init_key, warmup_key = jax.random.split(rng_key, 3)
init_keys = jax.random.split(init_key, n_chains)
init_params = jax.vmap(init_param_fn)(init_keys)

@jax.vmap
def call_warmup(seed, param):
    (initial_states, tuned_params), _ = warmup.run(seed, param, 1000)
    return initial_states, tuned_params

warmup_keys = jax.random.split(warmup_key, n_chains)
initial_states, tuned_params = call_warmup(warmup_keys, init_params)

CPU times: user 13.1 s, sys: 920 ms, total: 14 s
Wall time: 8.55 s

%%time
n_samples = 1000
rng_key, sample_key = jax.random.split(rng_key)
states, infos = inference_loop_multiple_chains(
    sample_key, initial_states, tuned_params, joint_logdensity_change_of_var, n_samples, n_chains
)

CPU times: user 11.9 s, sys: 415 ms, total: 12.4 s
Wall time: 4.15 s

# convert logits samples to theta samples
position = states.position._asdict()
position["thetas"] = jax.nn.sigmoid(position["logits"])
del position["logits"]  # delete logits
states = states._replace(position=position)

# make arviz trace from states
trace = arviz_trace_from_states(states, infos, burn_in=0)
summ_df = az.summary(trace)
summ_df
mean sd eti89_lb eti89_ub ess_bulk ess_tail r_hat mcse_mean mcse_sd
a 1 2 -2.9 3.6 7 4 1.54 1.1 0.64
b 11 6 2.9 22 7 4 1.53 2.4 1.8
thetas[0] 0.05 0.04 1.4e-09 0.13 7 4 1.53 0.014 0.012
thetas[1] 0.05 0.04 0.002 0.13 7 4 1.48 0.014 0.012
thetas[2] 0.05 0.04 5.7e-05 0.13 7 4 1.53 0.014 0.013
thetas[3] 0.05 0.05 2e-05 0.13 7 4 1.53 0.014 0.012
thetas[4] 0.05 0.05 3.4e-06 0.13 7 4 1.52 0.014 0.012
thetas[5] 0.05 0.04 1.3e-08 0.13 7 4 1.52 0.014 0.012
thetas[6] 0.05 0.05 0.0019 0.13 7 4 1.50 0.014 0.012
thetas[7] 0.05 0.05 1.6e-08 0.13 7 4 1.53 0.014 0.013
thetas[8] 0.05 0.05 2.1e-05 0.13 7 4 1.53 0.014 0.012
thetas[9] 0.05 0.05 2e-08 0.13 7 4 1.53 0.014 0.013
thetas[10] 0.05 0.05 0.00027 0.13 7 4 1.52 0.014 0.012
thetas[11] 0.05 0.05 6.4e-06 0.14 7 4 1.53 0.015 0.014
thetas[12] 0.05 0.05 0.00019 0.14 7 4 1.52 0.015 0.013
thetas[13] 0.05 0.05 3.7e-07 0.14 7 4 1.52 0.015 0.013
thetas[14] 0.08 0.05 0.029 0.17 13 5 1.20 0.013 0.012
thetas[15] 0.08 0.05 0.032 0.17 23 166 1.11 0.01 0.0092
thetas[16] 0.08 0.05 0.033 0.16 13 5 1.21 0.012 0.012
thetas[17] 0.11 0.05 0.03 0.17 18 144 1.14 0.013 0.0085
thetas[18] 0.09 0.05 0.032 0.18 112 147 1.28 0.0063 0.0065
thetas[19] 0.08 0.05 0.03 0.18 16 5 1.16 0.013 0.011
thetas[20] 0.09 0.05 0.033 0.18 79 160 1.19 0.0073 0.007
thetas[21] 0.08 0.06 0.016 0.18 8 4 1.41 0.018 0.014
thetas[22] 0.1 0.04 0.043 0.18 156 159 1.33 0.0054 0.0053
thetas[23] 0.107 0.042 0.046 0.18 5651 202 1.53 0.00051 0.00051
thetas[24] 0.12 0.05 0.044 0.19 42 167 1.06 0.0079 0.0059
thetas[25] 0.112 0.046 0.052 0.2 3899 196 1.52 0.0036 0.0036
thetas[26] 0.11 0.05 0.055 0.2 22 210 1.11 0.011 0.0093
thetas[27] 0.11 0.05 0.05 0.2 27 169 1.10 0.011 0.0097
thetas[28] 0.16 0.09 0.049 0.3 7 135 1.49 0.04 0.017
thetas[29] 0.1 0.06 0.043 0.21 11 4 1.25 0.017 0.013
thetas[30] 0.123 0.047 0.051 0.2 262 162 1.52 0.0018 0.0017
thetas[31] 0.11 0.06 0.044 0.23 81 209 1.17 0.0095 0.009
thetas[32] 0.116 0.034 0.06 0.17 154 186 1.39 0.0024 0.0021
thetas[33] 0.1 0.06 0.043 0.21 10 4 1.28 0.018 0.013
thetas[34] 0.121 0.035 0.064 0.18 143 139 1.32 0.0027 0.0023
thetas[35] 0.1 0.05 0.05 0.2 10 4 1.28 0.016 0.011
thetas[36] 0.11 0.06 0.041 0.22 9 4 1.35 0.02 0.015
thetas[37] 0.13 0.04 0.085 0.21 17 5 1.15 0.011 0.0094
thetas[38] 0.15 0.038 0.087 0.22 169 222 1.39 0.0029 0.0023
thetas[39] 0.12 0.08 0.023 0.24 7 4 1.52 0.028 0.017
thetas[40] 0.12 0.07 0.055 0.24 9 4 1.33 0.021 0.015
thetas[41] 0.143 0.058 0.067 0.25 4218 199 1.52 0.0041 0.0044
thetas[42] 0.19 0.05 0.11 0.25 31 149 1.08 0.0089 0.0063
thetas[43] 0.18 0.04 0.12 0.26 68 196 1.13 0.0068 0.0059
thetas[44] 0.16 0.06 0.093 0.28 17 5 1.15 0.016 0.013
thetas[45] 0.15 0.07 0.08 0.27 10 4 1.29 0.021 0.015
thetas[46] 0.19 0.06 0.093 0.27 70 174 1.12 0.007 0.0054
thetas[47] 0.15 0.07 0.055 0.27 7 4 1.47 0.027 0.016
thetas[48] 0.16 0.06 0.092 0.27 17 5 1.15 0.016 0.013
thetas[49] 0.15 0.07 0.069 0.27 9 4 1.36 0.023 0.016
thetas[50] 0.15 0.07 0.057 0.27 7 4 1.44 0.026 0.017
thetas[51] 0.17 0.06 0.11 0.27 9 4 1.31 0.018 0.012
thetas[52] 0.16 0.06 0.096 0.28 21 6 1.12 0.015 0.012
thetas[53] 0.17 0.06 0.093 0.29 47 191 1.08 0.012 0.01
thetas[54] 0.15 0.08 0.046 0.28 7 4 1.49 0.03 0.018
thetas[55] 0.17 0.07 0.093 0.3 10 4 1.29 0.022 0.015
thetas[56] 0.2 0.05 0.14 0.29 15 5 1.17 0.014 0.0096
thetas[57] 0.2 0.06 0.14 0.3 11 4 1.25 0.017 0.011
thetas[58] 0.202 0.059 0.11 0.31 7402 223 1.52 0.00067 0.00061
thetas[59] 0.18 0.07 0.1 0.31 10 4 1.30 0.022 0.016
thetas[60] 0.212 0.058 0.12 0.32 6206 201 1.52 0.00083 0.00075
thetas[61] 0.19 0.07 0.12 0.32 12 5 1.22 0.02 0.015
thetas[62] 0.22 0.059 0.13 0.33 6983 204 1.53 0.0007 0.00064
thetas[63] 0.22 0.06 0.14 0.34 206 183 1.42 0.0073 0.0064
thetas[64] 0.2 0.08 0.1 0.34 8 4 1.40 0.029 0.019
thetas[65] 0.19 0.09 0.07 0.34 7 4 1.50 0.035 0.02
thetas[66] 0.24 0.07 0.14 0.35 7 4 1.49 0.029 0.015
thetas[67] 0.25 0.07 0.17 0.37 8 4 1.40 0.025 0.014
thetas[68] 0.272 0.051 0.19 0.36 5989 149 1.53 0.00076 0.00065
thetas[69] 0.288 0.067 0.18 0.4 2931 176 1.52 0.0012 0.0011
thetas[70] 0.17 0.09 0.054 0.33 7 4 1.50 0.035 0.021

Hide code cell source

 
az.plot_trace(trace)
plt.tight_layout();
../_images/68ef7f13872fbca0e00da016c994b216a8dcfa976124a7f7d037099db0427825.png 
print(f"Number of divergence: {infos.is_divergent.sum()}")

Number of divergence: 0

We can see that r_hat is less than or equal to 1.01 for each latent variable, trace plots looks converged to stationary distribution, and only few samples are diverged.

Using a PPL#

Probabilistic programming language usually provides functionality to apply change of variable easily (often done automatically). In this case for TFP, we can use its modeling API tfd.JointDistribution*.

tfed = tfp.experimental.distributions

@tfd.JointDistributionCoroutineAutoBatched
def model():
    # TFP does not have improper prior, use uninformative prior instead
    a = yield tfd.HalfCauchy(0, 100, name='a')
    b = yield tfd.HalfCauchy(0, 100, name='b')
    yield tfed.IncrementLogProb(jnp.log(jnp.power(a + b, -2.5)), name='logdensity_ab')

    thetas = yield tfd.Sample(tfd.Beta(a, b), n_rat_tumors, name='thetas')
    yield tfd.Binomial(group_size, probs=thetas, name='y')

# Sample from the prior and prior predictive distributions. The result is a pytree.
# model.sample(seed=rng_key)

# Condition on the observed (and auxiliary variable).
pinned = model.experimental_pin(logdensity_ab=(), y=n_of_positives)
# Get the default change of variable bijectors from the model
bijectors = pinned.experimental_default_event_space_bijector()

rng_key, init_key = jax.random.split(rng_key)
prior_sample = pinned.sample_unpinned(seed=init_key)
# You can check the unbounded sample
# bijectors.inverse(prior_sample)

def joint_logdensity(unbound_param):
    param = bijectors.forward(unbound_param)
    log_det_jacobian = bijectors.forward_log_det_jacobian(unbound_param)
    return pinned.unnormalized_log_prob(param) + log_det_jacobian

%%time
warmup = blackjax.window_adaptation(blackjax.nuts, joint_logdensity)

# we use 4 chains for sampling
n_chains = 4
rng_key, init_key, warmup_key = jax.random.split(rng_key, 3)
init_params = bijectors.inverse(pinned.sample_unpinned(n_chains, seed=init_key))

@jax.vmap
def call_warmup(seed, param):
    (initial_states, tuned_params), _ = warmup.run(seed, param, 1000)
    return initial_states, tuned_params

warmup_keys = jax.random.split(warmup_key, n_chains)
initial_states, tuned_params = call_warmup(warmup_keys, init_params)

CPU times: user 17.7 s, sys: 1.16 s, total: 18.8 s
Wall time: 10.5 s

%%time
n_samples = 1000
rng_key, sample_key = jax.random.split(rng_key)
states, infos = inference_loop_multiple_chains(
    sample_key, initial_states, tuned_params, joint_logdensity, n_samples, n_chains
)

CPU times: user 10.4 s, sys: 510 ms, total: 10.9 s
Wall time: 4.69 s

# convert logits samples to theta samples
position = states.position
states = states._replace(position=bijectors.forward(position))

# make arviz trace from states
trace = arviz_trace_from_states(states, infos, burn_in=0)
summ_df = az.summary(trace)
summ_df
mean sd eti89_lb eti89_ub ess_bulk ess_tail r_hat mcse_mean mcse_sd
a 2.41 0.88 1.3 4 611 918 1.01 0.036 0.041
b 14.3 5.2 7.8 24 651 948 1.01 0.21 0.23
thetas[0] 0.064 0.042 0.012 0.14 2760 1840 1.00 0.00073 0.00072
thetas[1] 0.064 0.043 0.011 0.15 2602 1710 1.00 0.00075 0.00068
thetas[2] 0.064 0.041 0.013 0.14 3067 2355 1.00 0.00068 0.00068
thetas[3] 0.063 0.04 0.013 0.14 2688 2340 1.00 0.00071 0.00066
thetas[4] 0.064 0.042 0.012 0.14 2652 1996 1.00 0.00073 0.00065
thetas[5] 0.064 0.041 0.012 0.14 3058 2163 1.00 0.00065 0.0006
thetas[6] 0.063 0.042 0.01 0.14 2714 1899 1.00 0.00068 0.00067
thetas[7] 0.065 0.042 0.012 0.14 2917 2162 1.00 0.00069 0.00064
thetas[8] 0.066 0.042 0.013 0.14 2978 2193 1.00 0.00071 0.00076
thetas[9] 0.066 0.042 0.013 0.14 3046 2047 1.00 0.00068 0.0006
thetas[10] 0.066 0.043 0.012 0.14 2806 2068 1.00 0.00072 0.00066
thetas[11] 0.068 0.044 0.013 0.15 3265 2123 1.00 0.00069 0.00063
thetas[12] 0.069 0.046 0.012 0.15 2935 1993 1.00 0.00077 0.00069
thetas[13] 0.071 0.047 0.013 0.16 3681 2216 1.00 0.00071 0.00069
thetas[14] 0.091 0.048 0.027 0.18 5331 2323 1.00 0.00063 0.00052
thetas[15] 0.092 0.048 0.027 0.18 5032 2270 1.00 0.00061 0.0005
thetas[16] 0.093 0.049 0.028 0.18 4855 2521 1.00 0.00069 0.00064
thetas[17] 0.091 0.048 0.027 0.18 4392 2364 1.00 0.00072 0.00062
thetas[18] 0.094 0.049 0.029 0.19 4903 2483 1.00 0.00064 0.00056
thetas[19] 0.095 0.051 0.028 0.19 4928 2388 1.00 0.00069 0.00058
thetas[20] 0.097 0.05 0.03 0.19 4913 2423 1.00 0.00068 0.00054
thetas[21] 0.097 0.051 0.029 0.19 4595 2121 1.00 0.00072 0.00062
thetas[22] 0.105 0.0481 0.038 0.19 6312 2444 1.00 0.00058 0.00047
thetas[23] 0.107 0.05 0.039 0.2 5953 2692 1.00 0.00062 0.0005
thetas[24] 0.109 0.0499 0.041 0.2 6123 2427 1.00 0.00059 0.00048
thetas[25] 0.12 0.055 0.045 0.22 5895 2386 1.01 0.00069 0.00057
thetas[26] 0.119 0.053 0.047 0.21 5199 2413 1.00 0.0007 0.00057
thetas[27] 0.121 0.055 0.043 0.22 6417 2377 1.00 0.00066 0.00056
thetas[28] 0.119 0.054 0.044 0.21 5844 2339 1.00 0.00067 0.00055
thetas[29] 0.119 0.054 0.046 0.22 6473 2671 1.00 0.00066 0.00054
thetas[30] 0.119 0.054 0.044 0.21 6186 2499 1.00 0.00067 0.00055
thetas[31] 0.128 0.064 0.041 0.24 6562 2513 1.00 0.00077 0.00069
thetas[32] 0.1115 0.0381 0.058 0.18 5655 2795 1.00 0.00049 0.00037
thetas[33] 0.123 0.055 0.049 0.22 6264 2909 1.00 0.00066 0.0005
thetas[34] 0.117 0.0401 0.061 0.19 5810 2357 1.00 0.00051 0.00038
thetas[35] 0.124 0.0475 0.057 0.21 6236 2893 1.00 0.00059 0.00045
thetas[36] 0.131 0.058 0.051 0.23 6394 2714 1.00 0.00071 0.00055
thetas[37] 0.143 0.0418 0.082 0.21 6760 2886 1.00 0.0005 0.00037
thetas[38] 0.146 0.0452 0.08 0.22 5937 2814 1.00 0.00059 0.00046
thetas[39] 0.147 0.058 0.067 0.25 6652 2745 1.00 0.0007 0.00055
thetas[40] 0.147 0.058 0.066 0.25 6578 2435 1.00 0.00069 0.00054
thetas[41] 0.148 0.066 0.055 0.26 5534 2353 1.00 0.00086 0.0007
thetas[42] 0.177 0.0476 0.11 0.26 8091 3012 1.00 0.00053 0.0004
thetas[43] 0.185 0.0465 0.12 0.26 5904 2925 1.00 0.0006 0.00043
thetas[44] 0.175 0.065 0.083 0.29 5858 2680 1.00 0.00083 0.00065
thetas[45] 0.175 0.065 0.085 0.29 6225 2861 1.00 0.00081 0.00064
thetas[46] 0.176 0.063 0.087 0.29 7016 2319 1.00 0.00075 0.00064
thetas[47] 0.175 0.064 0.083 0.29 6336 2811 1.00 0.00079 0.00059
thetas[48] 0.176 0.063 0.085 0.28 7245 2961 1.00 0.00075 0.00056
thetas[49] 0.176 0.063 0.085 0.29 6918 2621 1.00 0.00076 0.00056
thetas[50] 0.175 0.064 0.084 0.29 6331 2890 1.00 0.0008 0.0006
thetas[51] 0.192 0.0494 0.12 0.27 6410 2833 1.00 0.00061 0.00048
thetas[52] 0.18 0.065 0.088 0.29 6636 2798 1.00 0.00078 0.00057
thetas[53] 0.181 0.063 0.09 0.29 6787 2578 1.00 0.00074 0.00055
thetas[54] 0.183 0.067 0.088 0.3 7400 2927 1.00 0.00079 0.00059
thetas[55] 0.192 0.063 0.1 0.3 6123 2803 1.00 0.00079 0.00056
thetas[56] 0.213 0.053 0.13 0.3 5323 2333 1.00 0.00071 0.00053
thetas[57] 0.22 0.0522 0.14 0.31 5675 2205 1.00 0.00068 0.00049
thetas[58] 0.203 0.068 0.11 0.32 6373 3003 1.00 0.00085 0.00065
thetas[59] 0.202 0.066 0.11 0.31 6027 2865 1.00 0.00083 0.00063
thetas[60] 0.212 0.064 0.12 0.32 5453 2782 1.00 0.00087 0.00067
thetas[61] 0.209 0.069 0.11 0.33 6314 2637 1.00 0.00087 0.00069
thetas[62] 0.219 0.066 0.12 0.33 6194 2948 1.00 0.00085 0.00063
thetas[63] 0.231 0.073 0.13 0.36 5189 2829 1.00 0.001 0.00074
thetas[64] 0.231 0.073 0.12 0.35 6795 2752 1.00 0.00088 0.00067
thetas[65] 0.231 0.073 0.12 0.36 6574 2882 1.00 0.0009 0.00073
thetas[66] 0.269 0.054 0.19 0.36 5329 2443 1.00 0.00072 0.00052
thetas[67] 0.279 0.059 0.19 0.38 4894 2794 1.00 0.00084 0.00058
thetas[68] 0.275 0.055 0.19 0.37 5130 2984 1.00 0.00078 0.00055
thetas[69] 0.284 0.071 0.18 0.4 4038 2771 1.00 0.0011 0.00082
thetas[70] 0.211 0.075 0.11 0.34 5460 3037 1.00 0.001 0.00074

Hide code cell source

 
az.plot_trace(trace)
plt.tight_layout();
../_images/b9436b89ca59453cdf1009d90605b9845b7874346444d34ea1aa8b3a3f41cbc1.png