inference gym / numpyro compatibility

Đang mở

#1,986 mở ngày 30/1/2025 1 bình luận 0 reaction 0 người được giao Xem trên GitHub

Chưa có ai nhận issue này.

Đánh giá

Độ khó: 4/5
Thời gian dự kiến: 3-5 ngày
Mức phù hợp với người mới: 25/100

Loại issue: Lỗi
Độ rõ ràng: Cần làm rõ
Mức độ hoạt động: Đình trệ
Công nghệ: jupyter-notebook, python
Lĩnh vực: machine-learning

Hướng nghiên cứu

Bắt đầu bằng cách tái hiện ví dụ trong notebook và đọc inference_gym/targets/banana.py quanh dòng 116, cùng với entry point inference_gym.using_jax được đề cập trong issue. Theo dõi cách độ cong được trace đi đến Banana target và xác định liệu đường dẫn JAX có hỗ trợ điều đó hay không; hoàn thành nghĩa là có một bản sửa khả năng tương thích đã được xác nhận hoặc một giới hạn được ghi rõ bằng tài liệu kèm theo một kiểm tra hồi quy.

Do mô hình lập chỉ mục viết ra từ nội dung của issue.

Mô tả

I'm looking to use Inference Gym targets in numpyro, but I'm running into issues I believe because there are numpy arrays in the Inference Gym model init, which causes tracer conversion errors in numpyro/jax

Any ideas how to get around this? I cant tell how the inference_gym.using_jax module works, but I was hoping that it would change the arrays to be initialized as jax arrays and not numpy

import jax
import numpyro
from numpyro import distributions as dist
from numpyro.infer import MCMC, NUTS
from inference_gym import using_jax as gym

class Banana(dist.Distribution):
    arg_constraints = {"ndims": dist.constraints.positive_integer, "curvature": dist.constraints.real}
    support = dist.constraints.real_vector
    pytree_data_fields = ("ndims", "curvature")

    def __init__(self, ndims, curvature):
        self.ndims = ndims
        self.curvature = curvature
        self.gym_dist = gym.targets.Banana(ndims=ndims, curvature=curvature)
        super().__init__(event_shape=(ndims,))
    
    def sample(self, key, sample_shape=()):
        return self.gym_dist.sample(seed=key, sample_shape=sample_shape)
    
    def log_prob(self, value):
        return self.gym_dist._unnormalized_log_prob(value)
    

samples = Banana(ndims=3, curvature=0.03).sample(jax.random.PRNGKey(0), sample_shape=(100,))

def model(X):
    curvature = numpyro.sample("curvature", dist.Beta(1,30))
    return numpyro.sample("obs", Banana(ndims=3, curvature=curvature), obs=X)

kernel = NUTS(model)
mcmc = MCMC(kernel, num_warmup=500, num_samples=1000)
mcmc.run(jax.random.PRNGKey(0), X=samples)

Here's the full traceback

Traceback

File "/var/folders/1p/v01fvg3j1cz1hzv988ygj8m00000gn/T/ipykernel_28885/3778384098.py", line 37, in
mcmc.run(jax.random.PRNGKey(0), X=samples)
File "/.venv/lib/python3.11/site-packages/numpyro/infer/mcmc.py", line 702, in run
states_flat, last_state = partial_map_fn(map_args)
^^^^^^^^^^^^^^^^^^^^^^^^
File "/.venv/lib/python3.11/site-packages/numpyro/infer/mcmc.py", line 465, in _single_chain_mcmc
new_init_state = self.sampler.init(
^^^^^^^^^^^^^^^^^^
File "/.venv/lib/python3.11/site-packages/numpyro/infer/hmc.py", line 749, in init
init_params = self._init_state(
^^^^^^^^^^^^^^^^^
File "/.venv/lib/python3.11/site-packages/numpyro/infer/hmc.py", line 693, in _init_state
) = initialize_model(
^^^^^^^^^^^^^^^^^
File "/.venv/lib/python3.11/site-packages/numpyro/infer/util.py", line 713, in initialize_model
(init_params, pe, grad), is_valid = find_valid_initial_params(
^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/.venv/lib/python3.11/site-packages/numpyro/infer/util.py", line 447, in find_valid_initial_params
(init_params, pe, z_grad), is_valid = _find_valid_params(
^^^^^^^^^^^^^^^^^^^
File "/.venv/lib/python3.11/site-packages/numpyro/infer/util.py", line 433, in _find_valid_params
_, _, (init_params, pe, z_grad), is_valid = init_state = body_fn(init_state)
^^^^^^^^^^^^^^^^^^^
File "/.venv/lib/python3.11/site-packages/numpyro/infer/util.py", line 417, in body_fn
pe, z_grad = value_and_grad(potential_fn)(params)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/.venv/lib/python3.11/site-packages/jax/_src/traceback_util.py", line 180, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^
File "/.venv/lib/python3.11/site-packages/jax/_src/api.py", line 468, in value_and_grad_f
ans, vjp_py = _vjp(f_partial, *dyn_args)
^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/.venv/lib/python3.11/site-packages/jax/_src/api.py", line 1975, in _vjp
out_primals, vjp = ad.vjp(flat_fun, primals_flat)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/.venv/lib/python3.11/site-packages/jax/_src/interpreters/ad.py", line 252, in vjp
out_primals, pvals, jaxpr, consts = linearize(traceable, *primals)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/.venv/lib/python3.11/site-packages/jax/_src/interpreters/ad.py", line 237, in linearize
jaxpr, out_pvals, consts = pe.trace_to_jaxpr_nounits(jvpfun_flat, in_pvals)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/.venv/lib/python3.11/site-packages/jax/_src/profiler.py", line 333, in wrapper
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/.venv/lib/python3.11/site-packages/jax/_src/interpreters/partial_eval.py", line 574, in trace_to_jaxpr_nounits
jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
^^^^^^^^^^^^^^^^^^^^^^^
File "/.venv/lib/python3.11/site-packages/jax/_src/linear_util.py", line 192, in call_wrapped
return self.f_transformed(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/.venv/lib/python3.11/site-packages/jax/_src/interpreters/partial_eval.py", line 587, in trace_to_subjaxpr_nounits
out_tracers, jaxpr, out_consts, env = _trace_to_subjaxpr_nounits(
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/.venv/lib/python3.11/site-packages/jax/_src/interpreters/partial_eval.py", line 616, in _trace_to_subjaxpr_nounits
ans = f(*in_args)
^^^^^^^^^^^
File "/.venv/lib/python3.11/site-packages/jax/_src/api_util.py", line 72, in flatten_fun
ans = f(*py_args, **py_kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^
File "/.venv/lib/python3.11/site-packages/jax/_src/interpreters/ad.py", line 78, in jvpfun
out_primals, out_tangents = f(tag, primals, tangents)
^^^^^^^^^^^^^^^^^^^^^^^^^
File "/.venv/lib/python3.11/site-packages/jax/_src/interpreters/ad.py", line 115, in jvp_subtrace
ans = f(*in_tracers)
^^^^^^^^^^^^^^
File "/.venv/lib/python3.11/site-packages/jax/_src/api_util.py", line 88, in flatten_fun_nokwargs
ans = f(*py_args)
^^^^^^^^^^^
File "/.venv/lib/python3.11/site-packages/jax/_src/api_util.py", line 292, in _argnums_partial
return fun(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/.venv/lib/python3.11/site-packages/numpyro/infer/util.py", line 299, in potential_energy
log_joint, model_trace = log_density(
^^^^^^^^^^^^^
File "/.venv/lib/python3.11/site-packages/numpyro/infer/util.py", line 70, in log_density
model_trace = trace(model).get_trace(*model_args, **model_kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/.venv/lib/python3.11/site-packages/numpyro/handlers.py", line 186, in get_trace
self(*args, **kwargs)
File "/.venv/lib/python3.11/site-packages/numpyro/primitives.py", line 105, in call
return self.fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^
File "/.venv/lib/python3.11/site-packages/numpyro/primitives.py", line 105, in call
return self.fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^
File "/.venv/lib/python3.11/site-packages/numpyro/primitives.py", line 105, in call
return self.fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^
[Previous line repeated 3 more times]
File "/var/folders/1p/v01fvg3j1cz1hzv988ygj8m00000gn/T/ipykernel_28885/3778384098.py", line 31, in model
return numpyro.sample("obs", Banana(ndims=3, curvature=curvature), obs=X)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/.venv/lib/python3.11/site-packages/numpyro/distributions/distribution.py", line 100, in call
return super().call(*args, kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/var/folders/1p/v01fvg3j1cz1hzv988ygj8m00000gn/T/ipykernel_28885/3778384098.py", line 17, in init
self.gym_dist = gym.targets.Banana(ndims=ndims, curvature=curvature)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/.venv/lib/python3.11/site-packages/inference_gym/targets/banana.py", line 116, in init
[10.] + [np.sqrt(1. + 2 * curvature2 * 10.**4)] +
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/.venv/lib/python3.11/site-packages/jax/_src/core.py", line 692, in array
raise TracerArrayConversionError(self)

Ngôn ngữ chính: Jupyter Notebook
Star: 4.4k
Fork: 1.1k
Chỉ số merge pull request: Không có pull request nào được merge trong 30 ngày

Chuẩn bị môi trường

Không có Dockerfile hay tệp Docker Compose
Không có mẫu pull request
Đọc hướng dẫn đóng góp

Bắt đầu từ đâu

Đọc hết issue, rồi đọc hướng dẫn đóng góp của dự án.
Bình luận trên issue rằng bạn sẽ nhận — tránh hai người làm cùng một việc.
Fork repository và làm thay đổi trên một nhánh.
Mở pull request có tham chiếu số hiệu của issue.

Issue khác của tensorflow/probability

KL divergence between two `TransformedDistribution`s with a shared non-injective bijector (`AbsoluteValue`) returns the base KL
Có thể đã có người làm @Andreas-Hum đã nhận 7 ngày trước. Đang mở
- Có gì không ổn?
Độ khó 2/5 1-3 giờ Mức phù hợp với người mới 72/100

tensorflow/probability#2036 · Mở ngày 1/11/2026
betaincinv skips the root finding when y == 0.5, so Beta medians are wrong
Đang mở
- Có gì không ổn?
Độ khó 3/5 1-2 ngày Mức phù hợp với người mới 68/100

tensorflow/probability#2034 · Mở ngày 8/10/2026
ENH: Implement ModAB root-finding algorithm
Đang mở
- Có gì không ổn?
Độ khó 5/5 Hơn một tuần Mức phù hợp với người mới 35/100

tensorflow/probability#2021 · Mở ngày 3/5/2026
Issues with transforming of a Mixture of Generalized Pareto using Bijector
Đang mở
- Có gì không ổn?
Độ khó 4/5 3-5 ngày Mức phù hợp với người mới 30/100

tensorflow/probability#2015 · Mở ngày 23/10/2025
PixelCNN example from the documentation does not work
Đang mở
- Có gì không ổn?
Độ khó 3/5 1-2 ngày Mức phù hợp với người mới 48/100

tensorflow/probability#2013 · Mở ngày 25/8/2025

Tất cả issue của tensorflow/probability

Issue tương tự

[Bug]: DoubleMLRDDData accepts a score_col that is no data column (as in its docstring example); the running variable then becomes a covariate and RDFlex fails
Có thể đã có người làm @vrjdev đã nhận hôm nay. Đang mở
- Có gì không ổn?
Độ khó 2/5 1-3 giờ Mức phù hợp với người mới 78/100

DoubleML/doubleml-for-py#455 · Mở ngày 9/11/2026

Maintainer thường phản hồi trong vòng 1 ngày
Install page: spectarf's CUDA engine fails with PyTorch for CUDA 13.0 (nvvmError: ERROR_INVALID_IR)
Có thể đã có người làm @lukepolson đã nhận hôm nay. Đang mở
- Có gì không ổn?
Độ khó 2/5 1-3 giờ Mức phù hợp với người mới 79/100

PyTomography/PyTomography#435 · Mở ngày 9/11/2026

Maintainer thường phản hồi trong vòng 1 ngày
Model.table(): Allow truncating of parameter list in distributions
Đang mở
- Có gì không ổn?
beginner friendly

Độ khó 2/5 1-3 giờ Mức phù hợp với người mới 67/100

pymc-devs/pymc#8609 · Mở ngày 9/11/2026

Maintainer thường phản hồi trong vòng 1 ngày
anthropic-claude/02 section 4 says the adaptive thinking page claims system and tool checkpoints survive a mode change; the page now says the opposite
Đang mở
- Có gì không ổn?
Độ khó 1/5 Dưới một giờ Mức phù hợp với người mới 82/100

aws-samples/sample-per-model-bedrock#821 · Mở ngày 9/11/2026

Maintainer thường phản hồi trong vòng 1 ngày
Give EThcD and ETciD their own fragmentation-type categories
Đang mở
- Có gì không ổn?
Độ khó 2/5 1-3 giờ Mức phù hợp với người mới 72/100

instadeepai/InstaNovo-FM#50 · Mở ngày 9/11/2026

Thêm issue về Machine Learning

inference gym / numpyro compatibility

Có gì không ổn?

Đánh giá

Hướng nghiên cứu

Mô tả

Chuẩn bị môi trường

Bắt đầu từ đâu

Issue khác của tensorflow/probability

Có gì không ổn?

Có gì không ổn?

Có gì không ổn?

Có gì không ổn?

Có gì không ổn?

Issue tương tự

Có gì không ổn?

Có gì không ổn?

Có gì không ổn?

Có gì không ổn?

Có gì không ổn?

Nhận issue mới trong hộp thư của bạn