Skip to content

Commit

Permalink
Switch jnp.reshape from newshape to shape paramter. (#20469)
Browse files Browse the repository at this point in the history
The newshape parameter was deprecated in JAX v0.4.28, and will soon be removed.
  • Loading branch information
jakevdp authored Nov 7, 2024
1 parent ccb07df commit 30a6b87
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion keras/src/backend/jax/image.py
Original file line number Diff line number Diff line change
Expand Up @@ -464,7 +464,7 @@ def affine_transform(
# transform the indices
coordinates = jnp.einsum("Bhwij, Bjk -> Bhwik", indices, transform)
coordinates = jnp.moveaxis(coordinates, source=-1, destination=1)
coordinates += jnp.reshape(a=offset, newshape=(*offset.shape, 1, 1, 1))
coordinates += jnp.reshape(a=offset, shape=(*offset.shape, 1, 1, 1))

# apply affine transformation
_map_coordinates = functools.partial(
Expand Down

0 comments on commit 30a6b87

Please sign in to comment.