(*args)
| 79 | |
| 80 | |
| 81 | def chain(*args): |
| 82 | args = [_to_partitioned_transformation(e) for e in args] |
| 83 | base = optax.chain(*[optax.GradientTransformation(init=e.init, update=e.update) for e in args]) |
| 84 | |
| 85 | def partition(param_spec): |
| 86 | return tuple(e.partition(param_spec) for e in args) |
| 87 | |
| 88 | return PartitionedGradientTransformation( |
| 89 | init=base.init, update=base.update, partition=partition |
| 90 | ) |
| 91 | |
| 92 | |
| 93 | def named_chain(**kwargs): |