From f45490c79a88977c70e64ba222a90ac96502d8ed Mon Sep 17 00:00:00 2001 From: Oleh Prypin Date: Wed, 14 Oct 2026 03:18:49 -0700 Subject: [PATCH] No public description PiperOrigin-RevId: 999251458 --- learned_optimization/tasks/parametric/image_mlp_ae.py | 2 +- learned_optimization/tasks/parametric/image_mlp_vae.py | 2 +- learned_optimization/tasks/parametric/lm_rnn.py | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/learned_optimization/tasks/parametric/image_mlp_ae.py b/learned_optimization/tasks/parametric/image_mlp_ae.py index c5026edb..b6494aa5 100644 --- a/learned_optimization/tasks/parametric/image_mlp_ae.py +++ b/learned_optimization/tasks/parametric/image_mlp_ae.py @@ -95,7 +95,7 @@ def init(self, key: PRNGKey) -> Params: image = next(self.datasets.train)["image"] # pyrefly: ignore[missing-attribute] return init_net(key, image) - def loss(self, params: Params, key: PRNGKey, data: Batch) -> jnp.ndarray: # pytype: disable=signature-mismatch # jax-ndarray + def loss(self, params: Params, key: PRNGKey, data: Batch) -> jnp.ndarray: # pyrefly: ignore[bad-override] unused_init_net, apply_net = hk.without_apply_rng( hk.transform(_forward)) diff --git a/learned_optimization/tasks/parametric/image_mlp_vae.py b/learned_optimization/tasks/parametric/image_mlp_vae.py index 4fc12b9b..bb26b638 100644 --- a/learned_optimization/tasks/parametric/image_mlp_vae.py +++ b/learned_optimization/tasks/parametric/image_mlp_vae.py @@ -126,7 +126,7 @@ def init(self, key: PRNGKey) -> Params: image = next(self.datasets.train)["image"] # pyrefly: ignore[missing-attribute] return hk.transform(_forward).init(key, image) - def loss(self, params: Params, key: PRNGKey, data: Batch) -> jnp.ndarray: # pytype: disable=signature-mismatch # jax-ndarray + def loss(self, params: Params, key: PRNGKey, data: Batch) -> jnp.ndarray: # pyrefly: ignore[bad-override] image = data["image"] center_fn = lambda img: 2 * (img - 0.5) image = jax.lax.cond(task_params["center_data"], center_fn, lambda x: x, diff --git a/learned_optimization/tasks/parametric/lm_rnn.py b/learned_optimization/tasks/parametric/lm_rnn.py index 8729001f..59372259 100644 --- a/learned_optimization/tasks/parametric/lm_rnn.py +++ b/learned_optimization/tasks/parametric/lm_rnn.py @@ -85,7 +85,7 @@ def _forward(inp): task_params.values["initializer"]) embed = hk.Embed(vocab_size, self.embed_size, w_init=w_init)(inp) - rnn_core = self.rnn_core_fn(self.rnn_size) # pytype: disable=wrong-arg-count # trace-all-classes + rnn_core = self.rnn_core_fn(self.rnn_size) # pyrefly: ignore[bad-argument-count] # Make learnable initial states. template_state = rnn_core.initial_state(1) @@ -124,7 +124,7 @@ def init(self, rng: PRNGKey) -> Params: seq = batch["obs"] return init_net(rng, seq) - def loss(self, params: Params, rng: PRNGKey, data: Batch) -> jnp.ndarray: # pytype: disable=signature-mismatch # jax-ndarray + def loss(self, params: Params, rng: PRNGKey, data: Batch) -> jnp.ndarray: net = hk.without_apply_rng(hk.transform(_forward)) obs = data["obs"] target = data["target"]