Function ComputeGradients

Computes gradients of variables in vmap w.r.t loss. Args: loss_or_activations: either the loss, which is a scalar tensor, or activations, which could be a tensor or a list of tensors. vmap: A `.NestedMap` of variables. grad_aggregation_method: Specifies the method used to combin

(
    loss_or_activations,
    vmap,
    grad_aggregation_method=tf.AggregationMethod.EXPERIMENTAL_TREE,
    colocate_gradients_with_ops=True,
    gate_gradients=False,
    compute_gradients_fn=None,
    skip_zero_gradients=None,
    use_bf16_gradients_ar=False,
    skip_none_gradients=True,
    defer_crs_to_apply_grad=False,
    activations_grad=None,
    is_activations=False,
    tpu_embedding_activations=None)