zea.backend.autograd

Autograd wrapper for different backends.

Classes

AutoGrad([verbose])

Wrapper class for autograd using different backends.

class zea.backend.autograd.AutoGrad(verbose=False)

Bases: object

Wrapper class for autograd using different backends.

property backend

Get Keras backend. Machine learning library of choice.

get_gradient_and_value_jit_fn(has_aux=False, disable_jit=False)

Returns a jitted function for calculating the gradients and function outputs.

get_gradient_jit_fn()

Returns a jitted function for calculating the gradients.

gradient(variable, **kwargs)

Returns the gradients of the function w.r.t. variable.

Parameters:

• variable (Tensor) – Input tensor.

• **kwargs – Keyword arguments to pass to self.function.

Returns:

Gradients of the function at variable.

∇f(x)

Return type:

gradients (Tensor)

gradient_and_value(variable, has_aux=False, **kwargs)

Returns both the gradients w.r.t. variable and outputs of the function.

Note that self.function should return a tuple of (out, aux) if has_aux=True. with aux being a tuple of auxiliary variables. If has_aux=False, self.function should return out only.

Parameters:

• variable (Tensor) – Input tensor.

• has_aux (bool) – Whether the function returns auxiliary variables.

• **kwargs – Keyword arguments to pass to self.function.

Returns:

Gradients of the function at variable.

∇f(x)

out (Tuple or Tensor): Outputs of the function at variable.

if has_aux: out = (f(x), aux) else: out = f(x)

Return type:

gradients (Tensor)

set_function(function)

Set the function to calculate the gradients of.