[LSC] Ignore incorrect type annotations related to jax.numpy APIs

Jake VanderPlas · The precondition Authors · commit e895bbff8f0c · 2023-09-26T07:11:41.000-07:00
PiperOrigin-RevId: 568529477
diff --git a/precondition/distributed_shampoo.py b/precondition/distributed_shampoo.py
@@ -1583,7 +1583,7 @@ def updated_statistics_from_grad(
       for axis in preconditioned_dims:
         update = functools.partial(gram_weighted_update, precision=precision)
         if frequent_directions:
-          if _should_compress(self._compression_rank, g.shape[axis]):
+          if _should_compress(self._compression_rank, g.shape[axis]):  # pytype: disable=wrong-arg-types  # jnp-type
             update = frequent_directions_update
         new_stat = update(to_float(stats[index]), g, axis, w1, w2)
         new_stats.append(from_float(new_stat))
diff --git a/precondition/tearfree/reallocation.py b/precondition/tearfree/reallocation.py
@@ -154,7 +154,7 @@ def score_fn(
     score_dict[name] = jnp.mean(
         jnp.array([ops_dict[rule](ct) for ct in current_target])
     )
-  return score_dict
+  return score_dict  # pytype: disable=bad-return-type  # jnp-type
 
 
 def create_redist_dict(

Original file line number	Diff line number	Diff line change
`@@ -154,7 +154,7 @@ def score_fn(`
`154`	`154`	`score_dict[name] = jnp.mean(`
`155`	`155`	`jnp.array([ops_dict[rule](ct) for ct in current_target])`
`156`	`156`	`)`
`157`		`- return score_dict`
	`157`	`+ return score_dict # pytype: disable=bad-return-type # jnp-type`
`158`	`158`
`159`	`159`
`160`	`160`	`def create_redist_dict(`