| 222 | self.assertEqual("contents 2", f.read()) |
| 223 | |
| 224 | def test_cond_prune(self, cycles): |
| 225 | x_in = [] |
| 226 | x_out = [] |
| 227 | |
| 228 | def f(x, y): |
| 229 | x_in.append(x) |
| 230 | xx = cond_v2.cond_v2( |
| 231 | math_ops.less(1, 2), |
| 232 | lambda: x + 1, |
| 233 | lambda: x + 2, |
| 234 | ) |
| 235 | x_out.append(xx) |
| 236 | return xx, 2 * y |
| 237 | |
| 238 | f_wrapped = wrap_function.wrap_function( |
| 239 | f, [tensor_spec.TensorSpec((), dtypes.float32)] * 2) |
| 240 | f_pruned = f_wrapped.prune(x_in[0], [x_out[0]]) |
| 241 | |
| 242 | class Adder(module.Module): |
| 243 | |
| 244 | @def_function.function(input_signature=[ |
| 245 | tensor_spec.TensorSpec(shape=None, dtype=dtypes.float32)]) |
| 246 | def add(self, x): |
| 247 | return f_pruned(x) |
| 248 | |
| 249 | root = Adder() |
| 250 | root.add(constant_op.constant(1.)) |
| 251 | root = cycle(root, cycles) |
| 252 | root.add(constant_op.constant(1.)) |
| 253 | |
| 254 | def test_capture_assets(self, cycles): |
| 255 | root = tracking.AutoTrackable() |