TypeError: The newshape argument to jnp.reshape was removed in JAX v0.4.36. Use shape instead.

[FFLLDDS]

Hello,

I get the error in the title when trying to dm_pix.affine_transform an image:

key = jax.random.key(0)
image = jax.random.normal(key=key, shape=(3, 32, 32)) 

matrix = jax.numpy.array([
        [1.0, 0.0, 0.0],
        [0.0, 1.0, 0.0],
        [0.0, 0.0, 1.0]
    ])

offset = jax.numpy.array([5, 5, 0]) 

img_t = dm_pix.affine_transform(image=image, matrix=matrix, offset=offset)

I have jax version 0.5.0 and dm_pix version 0.4.3 and installed with pip install dm-pix.

[pabloduque0]

Hey! Thanks for reaching out!

I am not able to reproduce currently.

Could you paste the full error trace and a pip freeze?

Thanks!

[FFLLDDS]

Hey, thank you!

...
img_t = dm_pix.affine_transform(image=image, matrix=matrix, offset=offset)
File "/opt/miniconda3/envs/repr/lib/python3.10/site-packages/dm_pix/_src/augment.py", line 718, in affine_transform
coordinates += jnp.reshape(a=offset, newshape=(*offset.shape, 1, 1, 1))
File "/opt/miniconda3/envs/repr/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py", line 2150, in reshape
raise TypeError("The newshape argument to jnp.reshape was removed in JAX v0.4.36."
TypeError: The newshape argument to jnp.reshape was removed in JAX v0.4.36. Use shape instead.

absl-py==2.1.0
chex==0.1.88
dm_pix==0.4.3
jax==0.5.0
jaxlib==0.5.0
ml_dtypes==0.5.1
numpy==2.2.2
opt_einsum==3.4.0
scipy==1.15.1
toolz==1.0.0
typing_extensions==4.12.2

[claudiofantacci]

Ah! This error was actually fixed in ad54e02, but I haven't done a release after it, so that's why you're getting this error!

You either need to install this package from main branch or wait a little for us to make a release. For this I just started the process to do it 😄

[claudiofantacci]

After #110 is submitted, I'll create a new release 😄

[claudiofantacci]

Hey @FFLLDDS I just released v0.4.4 and it's already available via pip. Can you please try again to see if the issue is resolved? 😄

[FFLLDDS]

Hey @claudiofantacci I updated to v0.4.4 and it works perfectly now! Thank you!

[claudiofantacci]

Fantastic! Thanks for the "wake up call" ahah 😄

[FFLLDDS]

No problem 😄 Thanks for the package! Really great having the JAX Ecosystem