Cannot run demo, possible incompatibility with latest Jax

[daniel-trejobanos]

Dear all,

I am trying to run the demo examples, but I run in the following error

---

ImportError Traceback (most recent call last) Input In [22], in <cell line: 1>() ----> 1 import bayesnewton 2 import objax 3 import numpy as np

File ~/opt/anaconda3/envs/aurora/lib/python3.9/site-packages/bayesnewton/init.py:1, in ----> 1 from . import ( 2 kernels, 3 utils, 4 ops, 5 likelihoods, 6 models, 7 basemodels, 8 inference, 9 cubature 10 ) 13 def build_model(model, inf, name='GPModel'): 14 return type(name, (inf, model), {})

File ~/opt/anaconda3/envs/aurora/lib/python3.9/site-packages/bayesnewton/kernels.py:5, in 3 import jax.numpy as np 4 from jax.scipy.linalg import cho_factor, cho_solve, block_diag, expm ----> 5 from jax.ops import index_add, index 6 from .utils import scaled_squared_euclid_dist, softplus, softplus_inv, rotation_matrix 7 from warnings import warn

ImportError: cannot import name 'index_add' from 'jax.ops' (/Users/Daniel/opt/anaconda3/envs/aurora/lib/python3.9/site-packages/jax/ops/init.py)

I think its related to this from the Jax website:

The functions jax.ops.index_update, jax.ops.index_add, etc., which were deprecated in JAX 0.2.22, have been removed. Please use the jax.numpy.ndarray.at property on JAX arrays instead.

[daniel-trejobanos]

I now realise that your pip installation asks for a specific jax version, which is a bit problematic for me, given that I am running on a M1 and installed jax via condaforge, I am not sure I can match to a compatible version, I will try and let you know if I succeed.

[daniel-trejobanos]

I managed to downgrade jax, but there is no jaxlib 0.1.60 available in condaforge, seems like it could be the source of this bug I get when trying to load objax 1.31:

---

TypeError Traceback (most recent call last) Input In [22], in <cell line: 1>() ----> 1 import bayesnewton 2 import objax 3 import numpy as np

File ~/opt/anaconda3/envs/aurora/lib/python3.9/site-packages/bayesnewton/init.py:1, in ----> 1 from . import ( 2 kernels, 3 utils, 4 ops, 5 likelihoods, 6 models, 7 basemodels, 8 inference, 9 cubature 10 ) 13 def build_model(model, inf, name='GPModel'): 14 return type(name, (inf, model), {})

File ~/opt/anaconda3/envs/aurora/lib/python3.9/site-packages/bayesnewton/kernels.py:1, in ----> 1 import objax 2 from jax import vmap 3 import jax.numpy as np

File ~/opt/anaconda3/envs/aurora/lib/python3.9/site-packages/objax/init.py:17, in 1 # Copyright 2020 Google LLC 2 # 3 # Licensed under the Apache License, Version 2.0 (the "License"); (...) 12 # See the License for the specific language governing permissions and 13 # limitations under the License. 15 import sys ---> 17 from ._patch_jax import * 19 pass # To avoid reordering imports from above 21 from . import functional

File ~/opt/anaconda3/envs/aurora/lib/python3.9/site-packages/objax/_patch_jax.py:20, in 16 all = [] 18 from typing import Union, Sequence, Tuple, Callable, Optional ---> 20 import jax.numpy as jn 22 from .typing import JaxArray 23 from .util import re_sign

File ~/opt/anaconda3/envs/aurora/lib/python3.9/site-packages/jax/init.py:93, in 89 from .version import version 91 # These submodules are separate because they are in an import cycle with 92 # jax and rely on the names imported above. ---> 93 from . import image 94 from . import lax 95 from . import nn

File ~/opt/anaconda3/envs/aurora/lib/python3.9/site-packages/jax/image/init.py:18, in 15 """Common functions for neural network libraries.""" 17 # flake8: noqa: F401 ---> 18 from jax._src.image.scale import ( 19 resize, 20 ResizeMethod, 21 scale_and_translate, 22 )

File ~/opt/anaconda3/envs/aurora/lib/python3.9/site-packages/jax/_src/image/scale.py:20, in 17 from typing import Callable, Sequence, Union 19 from jax import jit ---> 20 from jax import lax 21 from jax import numpy as jnp 22 import numpy as np

File ~/opt/anaconda3/envs/aurora/lib/python3.9/site-packages/jax/lax/init.py:324, in 291 from jax._src.lax.lax import (_reduce_sum, _reduce_max, _reduce_min, _reduce_or, 292 _reduce_and, _reduce_window_sum, _reduce_window_max, 293 _reduce_window_min, _reduce_window_prod, (...) 298 _upcast_fp16_for_computation, _broadcasting_shape_rule, 299 _eye, _tri, _delta, _ones, _zeros, _dilate_shape) 300 from jax._src.lax.control_flow import ( 301 associative_scan, 302 cond, (...) 322 while_p, 323 ) --> 324 from jax._src.lax.fft import ( 325 fft, 326 fft_p, 327 ) 328 from jax._src.lax.parallel import ( 329 all_gather, 330 all_to_all, (...) 346 xeinsum, 347 ) 348 from jax._src.lax.other import ( 349 conv_general_dilated_patches 350 )

File ~/opt/anaconda3/envs/aurora/lib/python3.9/site-packages/jax/_src/lax/fft.py:87, in 83 n = fft_lengths[-1] 84 return y[..., : n//2 + 1] 86 @partial(jit, static_argnums=1) ---> 87 def _rfft_transpose(t, fft_lengths): 88 # The transpose of RFFT can't be expressed only in terms of irfft. Instead of 89 # manually building up larger twiddle matrices (which would increase the 90 # asymptotic complexity and is also rather complicated), we rely JAX to 91 # transpose a naive RFFT implementation. 92 dummy_shape = t.shape[:-len(fft_lengths)] + fft_lengths 93 dummy_primal = ShapeDtypeStruct(dummy_shape, _real_dtype(t.dtype))

File ~/opt/anaconda3/envs/aurora/lib/python3.9/site-packages/jax/api.py:184, in jit(fun, static_argnums, device, backend, donate_argnums) 129 """Sets up fun for just-in-time compilation with XLA. 130 131 Args: (...) 181 -0.85743 -0.78232 0.76827 0.59566 ] 182 """ 183 if FLAGS.experimental_cpp_jit and config.omnistaging_enabled: --> 184 return _cpp_jit(fun, static_argnums, device, backend, donate_argnums) 185 else: 186 return _python_jit(fun, static_argnums, device, backend, donate_argnums)

File ~/opt/anaconda3/envs/aurora/lib/python3.9/site-packages/jax/api.py:370, in cpp_jit(fun, static_argnums, device, backend, donate_argnums) 367 return config.read("jax_disable_jit") 369 static_argnums = (0,) + tuple(i + 1 for i in static_argnums) --> 370 cpp_jitted_f = jax_jit.jit(fun, cache_miss, get_device_info, 371 get_jax_enable_x64, get_jax_disable_jit_flag, 372 static_argnums_) 374 # TODO(mattjj): make cpp callable follow descriptor protocol for bound methods 375 @wraps(fun) 376 @api_boundary 377 def f_jitted(*args, **kwargs):

TypeError: jit(): incompatible function arguments. The following argument types are supported: 1. (fun: function, cache_miss: function, get_device: function, static_argnums: List[int], static_argnames: List[str] = [], donate_argnums: List[int] = [], cache: jaxlib.xla_extension.CompiledFunctionCache = None) -> object

Invoked with: <function _rfft_transpose at 0x7f93913f20d0>, <function _cpp_jit..cache_miss at 0x7f93913f2160>, <function _cpp_jit..get_device_info at 0x7f93913f21f0>, <function _cpp_jit..get_jax_enable_x64 at 0x7f93913f2280>, <function _cpp_jit..get_jax_disable_jit_flag at 0x7f93913f2310>, (0, 2)

[wil-j-wil]

Hi,

Sorry for the slow response and apologies that you've been having issues with the package versions. This is indeed frustrating. I would love to update the package to work with the most recent versions, but I don't currently have the spare time.

I am using an M1 mac and things are working OK for me, but I'm not using condaforge.

The index_update issue should be fairly easy to fix. However, I recall seeing some performance issues when I tried updating objax in the past, and I never managed to debug the issue. I hope to get around to this at some point in the future.

[asolin]

This issue should be fixed here: https://github.com/AaltoML/BayesNewton/pull/18 (updated to current release of jax)