Getting a type error while using fori_loop with JAX Question: I’m developing a code using JAX, and I wanted to JIT some parts of that had big loops. I didn’t want the code to be unrolled so I used fori_loop, but I’m getting an error and can’t figure out what I am doing wrong. The …
Why is JAX's `split()` so slow at first call? Question: jax.numpy.split can be used to segment an array into equal-length segments with a remainder in the last element. e.g. splitting an array of 5000 elements into segments of 10: array = jnp.ones(5000) segment_size = 10 split_indices = jnp.arange(segment_size, array.shape[0], segment_size) segments = jnp.split(array, split_indices) This …
multidimensional jax.isin() Question: i am trying to filter an array of triples. The criterion by which I want to filter is whether another array of triples contains at least one element with the same first and third element. E.g import jax.numpy as jnp array1 = jnp.array( [ [0,1,2], [1,0,2], [0,3,3], [3,0,1], [0,1,1], [1,0,3], ] ) …
How to wrap a numpy function to make it work with jax.numpy? Question: I have some Jax code that requires using auto differentiation and in part of the code, I would like to call a function from a library written in NumPy. When I try this now I get The numpy.ndarray conversion method __array__() was …
Issues with non-hashable static arguments when forming Question: I have a vector-jacobian product that I want to compute. The function func takes four arguments, the final two of which are static: def func(variational_params, e, A, B): … return model_params, dlogp, … The function jits perfectly fine via func_jitted = jit(func, static_argnums=(2, 3)) The primals are …
Issue with jax.lax.scan Question: I am supposed to use Jax.lax.scan instead of a for loop with 100 iterations at line 22. I am supposed to update S and append it to S_list. I am unsure how to fix the jax.lax.scan. The error that keeps popping up is missing the required XS. When I put a …
jax segment_sum along array dimension Question: I am fairly new to jax and have the following problem: I need to compute functions (sum/min/max maybe more complex stuff later) across an array given an index. To solve this problem I found the jnp.ops.segment_sum function. This works great for one array, but how can I generalize this …
Error when trying to jit the computation of the Jacobian in JAX: "ValueError: Non-hashable static arguments are not supported" Question: This question is similar to the question here, but I cannot link with what I should alter. I have a function def elbo(variational_parameters, eps, a, b): … return theta, _ elbo = jit(elbo, static_argnames=["a", "b"]) …
Differences in the standard error of parameters computed via the hessian inverse and via QR decomposition Question: I have solved a nonlinear optimization and I am trying to compute the standard error of the parameters obtained. I found two options: One uses the fractional covariance matrix formed from the inverse of the hessian while the …
JAX – jitting functions: parameters vs "global" variables Question: I’ve have the following doubt about Jax. I’ll use an example from the official optax docs to illustrate it: def fit(params: optax.Params, optimizer: optax.GradientTransformation) -> optax.Params: opt_state = optimizer.init(params) @jax.jit def step(params, opt_state, batch, labels): loss_value, grads = jax.value_and_grad(loss)(params, batch, labels) updates, opt_state = optimizer.update(grads, opt_state, …
