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: 694438389
Change-Id: I75dbffb004999b235f3c36387ea4fd3e15e8c755

lightweight_mmm/optimize_media.py

@@ -71,7 +71,7 @@ def _objective_function(
   media_values = jnp.tile(
       media_values / media_input_shape[0], reps=media_input_shape[0])
   # Distribute budget of each channels across time.
-  media_values = jnp.reshape(a=media_values, newshape=media_input_shape)
+  media_values = jnp.reshape(a=media_values, shape=media_input_shape)
   media_values = media_scaler.transform(media_values)
   return -jnp.sum(
       media_mix_model.predict(