Compile the calculation in grad_fn if op was marked as compiled. Also apply mixed precision attributes to their respective gradient ops.
(scope, op, func, grad_fn)
| 333 | |
| 334 | |
| 335 | def _MaybeCompile(scope, op, func, grad_fn): |
| 336 | """Compile the calculation in grad_fn if op was marked as compiled. |
| 337 | Also apply mixed precision attributes to their respective gradient ops.""" |
| 338 | scope = scope.rstrip("/").replace("/", "_") |
| 339 | if func is not None: |
| 340 | xla_compile = func.definition.attr["_XlaCompile"].b |
| 341 | xla_separate_compiled_gradients = func.definition.attr[ |
| 342 | "_XlaSeparateCompiledGradients"].b |
| 343 | xla_scope = func.definition.attr["_XlaScope"].s.decode() |
| 344 | # TODO(mconley) Add mixed precision scope support to funcdefs |
| 345 | auto_mixed_precision_scope = False |
| 346 | else: |
| 347 | try: |
| 348 | auto_mixed_precision_include = op.get_attr( |
| 349 | "_AutoMixedPrecisionScopeInclude") |
| 350 | auto_mixed_precision_scope = True |
| 351 | except ValueError: |
| 352 | auto_mixed_precision_scope = False |
| 353 | try: |
| 354 | xla_compile = op.get_attr("_XlaCompile") |
| 355 | xla_separate_compiled_gradients = op.get_attr( |
| 356 | "_XlaSeparateCompiledGradients") |
| 357 | xla_scope = op.get_attr("_XlaScope").decode() |
| 358 | except ValueError: |
| 359 | xla_compile = False |
| 360 | |
| 361 | if not (xla_compile or auto_mixed_precision_scope): |
| 362 | return grad_fn() # Exit early |
| 363 | |
| 364 | if auto_mixed_precision_scope: |
| 365 | auto_mixed_precision_attrs = { |
| 366 | "_AutoMixedPrecisionScopeInclude": attr_value_pb2.AttrValue( |
| 367 | b=auto_mixed_precision_include) |
| 368 | } |
| 369 | else: |
| 370 | auto_mixed_precision_attrs = {} |
| 371 | |
| 372 | if not xla_compile: |
| 373 | with ops.get_default_graph()._attr_scope( # pylint: disable=protected-access |
| 374 | auto_mixed_precision_attrs): |
| 375 | return grad_fn() |
| 376 | |
| 377 | |
| 378 | # If the gradients are supposed to be compiled separately, we give them a |
| 379 | # _XlaScope name that is based on the name_scope of the gradients. Otherwise |
| 380 | # they just inherit the existing _XlaScope name, which lets them be merged |
| 381 | # together with the non-gradient computation. |
| 382 | if xla_separate_compiled_gradients: |
| 383 | xla_grad_scope = "%s_grad_%s" % (xla_scope, scope) |
| 384 | else: |
| 385 | xla_grad_scope = xla_scope |
| 386 | |
| 387 | attrs = { |
| 388 | "_XlaCompile": attr_value_pb2.AttrValue(b=xla_compile), |
| 389 | "_XlaScope": attr_value_pb2.AttrValue(s=xla_grad_scope.encode()) |
| 390 | } |
| 391 | attrs.update(auto_mixed_precision_attrs) |
| 392 | with ops.get_default_graph()._attr_scope(attrs): # pylint: disable=protected-access |
no test coverage detected