Skip to content

fix: fix transpose empty iterable with zip(*nested) in transformations #145

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 10 commits into from
Mar 7, 2023
Prev Previous commit
Next Next commit
style: use postional-only arguments
  • Loading branch information
XuehaiPan committed Mar 3, 2023
commit 95af6b1e079630cdae8ce14fe19a95d4afcffa89
35 changes: 19 additions & 16 deletions torchopt/diff/implicit/nn/module.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,38 +36,40 @@


def _stateless_objective_fn(
__flat_params: TupleOfTensors,
__flat_meta_params: TupleOfTensors,
__params_names: Iterable[str],
__meta_params_names: Iterable[str],
flat_params: TupleOfTensors,
flat_meta_params: TupleOfTensors,
params_names: Iterable[str],
meta_params_names: Iterable[str],
self: ImplicitMetaGradientModule,
/,
*input: Any,
**kwargs: Any,
) -> torch.Tensor:
with reparametrize(
self,
itertools.chain(
zip(__params_names, __flat_params),
zip(__meta_params_names, __flat_meta_params),
zip(params_names, flat_params),
zip(meta_params_names, flat_meta_params),
),
):
return self.objective(*input, **kwargs)


def _stateless_optimality_fn(
__flat_params: TupleOfTensors,
__flat_meta_params: TupleOfTensors,
__params_names: Iterable[str],
__meta_params_names: Iterable[str],
flat_params: TupleOfTensors,
flat_meta_params: TupleOfTensors,
params_names: Iterable[str],
meta_params_names: Iterable[str],
self: ImplicitMetaGradientModule,
/,
*input: Any,
**kwargs: Any,
) -> TupleOfTensors:
with reparametrize(
self,
itertools.chain(
zip(__params_names, __flat_params),
zip(__meta_params_names, __flat_meta_params),
zip(params_names, flat_params),
zip(meta_params_names, flat_meta_params),
),
):
return self.optimality(*input, **kwargs)
Expand Down Expand Up @@ -121,12 +123,13 @@ def enable_implicit_gradients(
@custom_root(_stateless_optimality_fn, argnums=1, has_aux=True, **solve_kwargs)
def stateless_solver_fn(
# pylint: disable=unused-argument
__flat_params: TupleOfTensors,
__flat_meta_params: TupleOfTensors,
__params_names: Iterable[str],
__meta_params_names: Iterable[str],
flat_params: TupleOfTensors,
flat_meta_params: TupleOfTensors,
params_names: Iterable[str],
meta_params_names: Iterable[str],
# pylint: enable=unused-argument
self: ImplicitMetaGradientModule,
/,
*input: Any,
**kwargs: Any,
) -> tuple[TupleOfTensors, Any]:
Expand Down
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy