Returns a function that computes f and its vjp w.r.t. params. The term "vjp" here is an abbreviation for vector-jacobian product. Args: f: the function to be differentiated. params: the parameters (numbers or names) to differentiate with respect to. A value of None will differ
(f, params=None, persistent=True)
| 472 | |
| 473 | |
| 474 | def make_vjp(f, params=None, persistent=True): |
| 475 | """Returns a function that computes f and its vjp w.r.t. |
| 476 | |
| 477 | params. |
| 478 | |
| 479 | The term "vjp" here is an abbreviation for vector-jacobian product. |
| 480 | |
| 481 | Args: |
| 482 | f: the function to be differentiated. |
| 483 | params: the parameters (numbers or names) to differentiate with respect to. |
| 484 | A value of None will differentiate with respect to all parameters. |
| 485 | persistent: Boolean controlling whether the VJP function can be re-used. |
| 486 | Must be True or False. |
| 487 | |
| 488 | Returns: |
| 489 | A function, which when called, returns a tuple (value, vjp), where: |
| 490 | - value is the result of calling f. |
| 491 | - vjp is a function, which takes a vector as an argument and |
| 492 | returns the product of that vector with the Jacobian of f. |
| 493 | Providing no argument to vjp is equivalent to providing a |
| 494 | vector of ones. |
| 495 | |
| 496 | For example, |
| 497 | ```python |
| 498 | def f(x): |
| 499 | return x * x |
| 500 | |
| 501 | wrapped_fn = tfe.make_vjp(f) |
| 502 | result, vjp = wrapped_fn(tf.constant(3.0)) |
| 503 | # result is 9.0 |
| 504 | vjp() # the vjp function rturns 6.0 |
| 505 | |
| 506 | Raises: |
| 507 | ValueError: if `f` returns None. |
| 508 | """ |
| 509 | |
| 510 | def decorated(*args, **kwds): |
| 511 | """Computes the value and gradient of the decorated function.""" |
| 512 | parameter_positions = _get_arg_spec(f, params, args) |
| 513 | assert not kwds, "The gradient function can't take keyword arguments." |
| 514 | this_tape = tape.push_new_tape(persistent=persistent) |
| 515 | try: |
| 516 | sources = [] |
| 517 | args = [ |
| 518 | ops.convert_to_tensor(arg) if i in parameter_positions else arg |
| 519 | for i, arg in enumerate(args) |
| 520 | ] |
| 521 | args = _ensure_unique_tensor_objects(parameter_positions, args) |
| 522 | for i in parameter_positions: |
| 523 | sources.append(args[i]) |
| 524 | tape.watch(this_tape, args[i]) |
| 525 | result = f(*args) |
| 526 | if result is None: |
| 527 | raise ValueError("Cannot differentiate a function that returns None; " |
| 528 | "did you forget to return a value from {}?".format( |
| 529 | f.__name__)) |
| 530 | flat_result = nest.flatten(result) |
| 531 | flat_result = [gen_array_ops.identity(x) for x in flat_result] |