Silence some pytype errors related to a JAX build refactor

This build change allows pytype to propagate annotations that it previously did not, and because of this it starts flagging existing incorrect annotations.

PiperOrigin-RevId: 772480864
Change-Id: If8bd74414306e120764eaf0f9eabe725d5282680

Author: Jake VanderPlas

lightweight_mmm/core/transformations/lagging.py

@@ -152,7 +152,7 @@ def adstock_internal(
     adstock_value = prev_adstock * lag_weight + data
     return adstock_value, adstock_value# jax-ndarray
 
-  _, adstock_values = jax.lax.scan(
+  _, adstock_values = jax.lax.scan(# lax-types
       f=adstock_internal, init=data[0, ...], xs=data[1:, ...])
   adstock_values = jnp.concatenate([jnp.array([data[0, ...]]), adstock_values])
   return jax.lax.cond(

lightweight_mmm/media_transforms.py

@@ -84,7 +84,7 @@ def adstock_internal(prev_adstock: jnp.ndarray,
     adstock_value = prev_adstock * lag_weight + data
     return adstock_value, adstock_value# jax-ndarray
 
-  _, adstock_values = jax.lax.scan(
+  _, adstock_values = jax.lax.scan(# lax-types
       f=adstock_internal, init=data[0, ...], xs=data[1:, ...])
   adstock_values = jnp.concatenate([jnp.array([data[0, ...]]), adstock_values])
   return jax.lax.cond(