In jnp.reshape, use shape rather than deprecated newshape. The `n…

…ewshape` parameter was deprecated in JAX v0.4.28, and will soon be removed. PiperOrigin-RevId: 694193311
google-deepmind · Nov 8, 2024 · 261527d · 261527d
1 parent b76b0b6
commit 261527d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/dm_pix/_src/augment.py b/dm_pix/_src/augment.py
@@ -715,7 +715,7 @@ def affine_transform(
 
   # Alter coordinates to account for offset.
   offset = jnp.full((3,), fill_value=offset)
-  coordinates += jnp.reshape(a=offset, newshape=(*offset.shape, 1, 1, 1))
+  coordinates += jnp.reshape(a=offset, shape=(*offset.shape, 1, 1, 1))
 
   interpolate_function = _get_interpolate_function(
       mode=mode,