parameters: (op_name,
                    attr_tuple,
                    inputs,
                    outputs,
                    tangents,
                    use_batch=False)
return_statement: return _jvp_relaxed_shapes(op_name, attr_tuple, inputs, outputs, tangents,
                               use_batch)
docstring: Determine which forwardprop function to call.
docstring_summary: Determine which forwardprop function to call.
docstring_tokens: Determine which forwardprop function to call .
function: def _jvp_dispatch(op_name,
                    attr_tuple,
                    inputs,
                    outputs,
                    tangents,
                    use_batch=False):
    """Determine which forwardprop function to call."""
    # Note that this _TRACE_COUNT read races with writes. That's fine, it just
    # means we may trace a few more exact shapes before moving on to relaxation.
    if _TRACE_COUNT.get(op_name, 0) < _TRACE_COUNT_LIMIT:
      return _jvp_exact_shapes(op_name, attr_tuple, inputs, outputs, tangents,
                               use_batch)
    return _jvp_relaxed_shapes(op_name, attr_tuple, inputs, outputs, tangents,
                               use_batch)
function_tokens: def _jvp_dispatch ( op_name , attr_tuple , inputs , outputs , tangents , use_batch = False ) : # Note that this _TRACE_COUNT read races with writes. That's fine, it just # means we may trace a few more exact shapes before moving on to relaxation. if _TRACE_COUNT . get ( op_name , 0 ) < _TRACE_COUNT_LIMIT : return _jvp_exact_shapes ( op_name , attr_tuple , inputs , outputs , tangents , use_batch ) return _jvp_relaxed_shapes ( op_name , attr_tuple , inputs , outputs , tangents , use_batch )