in blocksparse/grads.py [0:0]
def gradients(ys, xs, grad_ys=None, stop_grads=None, group_aggregations=8, custom_matmul_grad=True):
if group_aggregations > 8 or group_aggregations < 1:
raise ValueError("gradients: group_aggregation sizes of 1-8 supported.")
ys = _AsList(ys)
xs = [x.value() if isinstance(x, tf.Variable) else x for x in _AsList(xs)]
stop_grads = [] if stop_grads is None else _AsList(stop_grads)
grad_ys = [None] * len(ys) if grad_ys is None else _AsList(grad_ys)
assert len(ys) == len(grad_ys)
with ops.name_scope("gradients"):
for i, dy in enumerate(grad_ys):
if dy is None:
# float grads start at ones by default
grad_ys[i] = tf.fill(tf.shape(ys[i]), tf.constant(1.0, dtype=ys[i].dtype, name=f"grad_ys_{i}"))
ys_ops = [t.op for t in ys]
xs_ops = [t.op for t in xs]
pending_count, reachable_ys_ops, recompute_ops = _PendingCount(ys_ops, xs_ops)
# The set of ops that terminate the gradient computation.
# Confirm that our xs tensors are just endpoints in the graph.
# Also set any externally provided stop grad ops.
stop_ops = set(t.op for t in stop_grads)
for op in xs_ops:
is_stop_op = True
for x in op.inputs:
if x.op in pending_count:
is_stop_op = False
break
if is_stop_op:
stop_ops.add(op)
# Each op output has an associated list of gradient inputs
# If more than one, these need to be accumulated.
# Add the initial gradients for the ys.
grads = dict()
for y, dy in zip(ys, grad_ys):
_SetGrad(grads, y, dy)
# Add the unique ys ops that are ready into the queue.
queue = collections.deque()
for op in reachable_ys_ops:
# an op is ready if it has no dependecies
if op not in pending_count:
queue.append(op)
while queue:
op = queue.popleft()
# only pending_count==0 ops are in the queue so all grad input lists are fully populated
# go ahead and apply any needed add_n ops to these lists.
dys = _AggregatedGrads(grads, op, group_aggregations)
# confirm that we have at least one tensor to compute and that this isn't a stop grad op
if any(dy is not None for dy in dys) and op not in stop_ops:
# get the grad function for this op
try:
if custom_matmul_grad and op.type == "MatMul" and not op.get_attr("transpose_a") and not op.get_attr("transpose_b"):
grad_fn = _MatMulGradNN
else:
grad_fn = ops.get_gradient_function(op)
except LookupError:
raise LookupError(f"No gradient defined for operation '{op.name}' (op type: {op.type})")
# for any missing input grads, build a zero input of the right dtype/shape
for i, dy in enumerate(dys):
if dy is None:
dys[i] = tf.zeros_like(op.outputs[i])
# call the grad function with the forward op node and list of grad inputs
with ops.name_scope(op.name + "_grad"):
dxs = _AsList(grad_fn(op, *dys))
if len(dxs) != len(op.inputs):
raise ValueError(f"Num gradients {len(dxs)} generated for op {op.node_def} do not match num inputs {len(op.inputs)}")
#_LogOpGradients(op, dys, dxs)
else:
dxs = [None] * len(op.inputs)
for i, (x, dx) in enumerate(zip(op.inputs, dxs)):
if dx is not None:
# force unsorted_segment_sum call
if isinstance(dx, ops.IndexedSlices):
dx = tf.convert_to_tensor(dx)
#dx = emb.embedding_lookup_grad_op(dx.values, dx.indices, dx.dense_shape[0])
# do some shape sanity checking
try:
dx.set_shape(x.shape)
except ValueError:
raise ValueError("Incompatible shapes between op input {x.shape} and calculated input gradient {dx.shape} for {op.name} (idx:{i})")
# update the input grad list for the consumer of this gradient
_SetGrad(grads, x, dx)
# Update pending count for the inputs of op and enqueue any ready ops
for x in op.inputs:
# only traverse nodes that are in the reachable gradient path (and hence have a pending entry)
count = pending_count.get(x.op)
if count is not None:
if count == 1:
# when count is 1 this should be last time we reach this node
queue.append(x.op)
pending_count[x.op] = count - 1
# Disconnect the recomputed portion of the graph from the forward pass.
# This was only needed to direct the gradient flow.
# Leaving these connections in place would create a circular dependancy (from added control inputs).
for op in recompute_ops:
# Just overwrite the backward inputs with a copy of the forward inputs.
n_out = len(op.outputs)
for i, x in enumerate(op.inputs[:n_out]):
op._update_input(i+n_out, x)
return [_GetGrad(grads, x) for x in xs]